fix ci nn 5

EnzymeAD · Jan 31, 2025 · 08e800e · 08e800e
1 parent 379debc
commit 08e800e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/ext/ReactantNNlibExt.jl b/ext/ReactantNNlibExt.jl
@@ -175,8 +175,10 @@ function reduce_window(f, x::AnyTracedRArray{T,N}, pdims; init) where {T,N}
             body
         end
 
-    attr = fill(Reactant.MLIR.IR.Attribute(init), unranked)
-    init_value = Reactant.MLIR.IR.result(Ops.constant(; value=attr))
+    attr = Reactant.MLIR.IR.DenseElements(fill(Reactant.MLIR.IR.Attribute(init), unranked))
+    init_value = Reactant.MLIR.IR.result(
+        Reactant.MLIR.Dialects.stablehlo.constant(; value=attr)
+    )
     reduction = Reactant.MLIR.Dialects.stablehlo.reduce_window(
         [get_mlir_data(x)],
         [init_value];

diff --git a/src/mlir/IR/Attribute.jl b/src/mlir/IR/Attribute.jl
@@ -873,6 +873,7 @@ Base.show(io::IO, f::FlatSymbol) = print(io, "@$(flatsymbol(f.attr))")
 struct DenseElements{T}
     attr::Attribute
     DenseElements(a::AbstractArray{T}) where {T} = new{T}(IR.DenseElementsAttribute(a))
+    DenseElements(a::Attribute) = begin if !isdenseelements(a) throw("$a is not a dense elements attribute.") end; new{:Generic}(a) end
 end
 Attribute(d::DenseElements) = d.attr