fix: bypass segfault with fill complex (#245)

* fix: bypass segfault with fill complex

* test: simpler promote_to testing

* fix: `MLIR.IR.DenseElementsAttribute` definition for bool arrays

src/TracedRArray.jl

@@ -697,8 +697,7 @@ function broadcast_to_size(arg::Base.RefValue, rsize)
 end
 
 function broadcast_to_size(arg::T, rsize) where {T<:Number}
-    TT = MLIR.IR.TensorType([Int64(s) for s in rsize], MLIR.IR.Type(typeof(arg)))
-    attr = Base.fill(arg, TT)
+    attr = MLIR.IR.DenseElementsAttribute(Base.fill(arg, Tuple(rsize)))
     return arg = TracedRArray{T,length(rsize)}(
         (), MLIR.IR.result(MLIR.Dialects.stablehlo.constant(; value=attr), 1), rsize
     )

src/TracedRNumber.jl

@@ -86,7 +86,7 @@ function promote_to(::Type{TracedRNumber{T}}, rhs) where {T}
         )
     end
     if isa(rhs, Number)
-        attr = fill(MLIR.IR.Attribute(T(rhs)), mlir_type(TracedRNumber{T}))
+        attr = MLIR.IR.DenseElementsAttribute(fill(T(rhs)))
         return TracedRNumber{T}(
             (), MLIR.IR.result(MLIR.Dialects.stablehlo.constant(; value=attr), 1)
         )

src/mlir/IR/Attribute.jl

@@ -501,7 +501,7 @@ end
 
 Creates a dense elements attribute with the given shaped type from elements of a specific type. Expects the element type of the shaped type to match the data element type.
 """
-function DenseElementsAttribute(values::AbstractVector{Bool})
+function DenseElementsAttribute(values::AbstractArray{Bool})
     shaped_type = TensorType(size(values), Type(Bool))
     return Attribute(
         API.mlirDenseElementsAttrBoolGet(shaped_type, length(values), pointer(values))

test/complex.jl

@@ -86,3 +86,14 @@ end
     x_concrete = Reactant.to_rarray(x)
     @test @jit(abs.(x_concrete)) ≈ abs.(x)
 end
+
+@testset "promote_to Complex" begin
+    x = 1.0 + 2.0im
+    y = Reactant.ConcreteRNumber(x)
+
+    f = Reactant.compile((y,)) do z
+        z + Reactant.promote_to(Reactant.TracedRNumber{ComplexF64}, 1.0 - 3.0im)
+    end
+
+    @test isapprox(f(y), 2.0 - 1.0im)
+end

test/nn/nnlib.jl

@@ -157,4 +157,9 @@ end
 
     pad_fn2 = Base.Fix2(NNlib.pad_constant, (1, 0, 1, 3))
     @test @jit(∇sumabs2(pad_fn2, x_ra)) ≈ ∇sumabs2(pad_fn2, x)
+
+    x = rand(ComplexF32, 4, 4)
+    x_ra = Reactant.ConcreteRArray(x)
+
+    @test @jit(NNlib.pad_constant(x_ra, (1, 1))) ≈ NNlib.pad_constant(x, (1, 1))
 end