feat: add Base.Complex overloads (#1385)

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.129"
+version = "0.2.130"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/TracedRNumber.jl

@@ -119,6 +119,22 @@ function TracedUtils.promote_to(::TracedRNumber{T}, rhs) where {T}
     return TracedUtils.promote_to(TracedRNumber{T}, rhs)
 end
 
+for (aT, bT) in (
+    (TracedRNumber{<:Real}, Real),
+    (Real, TracedRNumber{<:Real}),
+    (TracedRNumber{<:Real}, TracedRNumber{<:Real}),
+)
+    @eval function Base.Complex(a::$aT, b::$bT)
+        T = promote_type(unwrapped_eltype(a), unwrapped_eltype(b))
+        a = TracedUtils.promote_to(TracedRNumber{T}, a)
+        b = TracedUtils.promote_to(TracedRNumber{T}, b)
+        return Ops.complex(a, b)
+    end
+end
+
+Base.Complex(x::TracedRNumber{<:Real}) = Ops.complex(x, zero(x))
+Base.Complex(x::TracedRNumber{<:Complex}) = x
+
 for (jlop, hloop) in (
     (:(Base.min), :minimum),
     (:(Base.max), :maximum),

test/complex.jl

@@ -103,3 +103,18 @@ end
     x_ra = Reactant.to_rarray(x)
     @test @jit(sum(abs2, x_ra)) ≈ sum(abs2, x)
 end
+
+@testset "create complex numbers" begin
+    x = randn(ComplexF32)
+    x_ra = Reactant.to_rarray(x; track_numbers=true)
+    @test @jit(Complex(x_ra)) == x_ra
+
+    x = randn(Float32)
+    y = randn(Float64)
+    x_ra = Reactant.to_rarray(x; track_numbers=true)
+    y_ra = Reactant.to_rarray(y; track_numbers=true)
+    @test @jit(Complex(x_ra, y_ra)) == Complex(x, y)
+    @test @jit(Complex(x_ra, y)) == Complex(x, y)
+    @test @jit(Complex(x, y_ra)) == Complex(x, y)
+    @test @jit(Complex(x_ra)) == Complex(x) == @jit(Complex(x_ra, 0))
+end