fix abs of complex number (#239)

* fix abs of complex number * Update complex.jl * Update complex.jl
EnzymeAD · Nov 7, 2024 · c4d0b50 · c4d0b50
1 parent 06f2c36
commit c4d0b50
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/src/TracedRNumber.jl b/src/TracedRNumber.jl
@@ -228,7 +228,8 @@ for (jlop, hloop) in (
     (:(Base.sqrt), :sqrt),
 )
     @eval function $(jlop)(@nospecialize(lhs::TracedRNumber{T})) where {T}
-        return TracedRNumber{T}(
+        OutTy = $(hloop === :abs) ? real(T) : T
+        return TracedRNumber{OutTy}(
             (), MLIR.IR.result(MLIR.Dialects.stablehlo.$hloop(lhs.mlir_data), 1)
         )
     end

diff --git a/test/complex.jl b/test/complex.jl
@@ -80,3 +80,9 @@ end
         @test @jit(imag(x_concrete)) == imag(x)
     end
 end
+
+@testset "abs: $T" for T in (Float32, ComplexF32)
+    x = randn(T, 10)
+    x_concrete = Reactant.to_rarray(x)
+    @test @jit(abs.(x_concrete)) ≈ abs.(x)
+end