Support multiplication by im (#1929)

* multiplication by im

* test

* fix

* Float64 --> Float32

Author: jumerckx

src/TracedRNumber.jl

@@ -242,6 +242,15 @@ for (jlop, hloop) in (
     end
 end
 
+function Base.:*(x::TracedRNumber{T}, z::Complex{Bool}) where {T<:Real}
+    # this is to support multiplication by im (Complex{Bool}(false, true))
+    z_re, z_im = real(z), imag(z)
+    res_re = z_re ? x : zero(x)
+    res_im = z_im ? x : zero(x)
+    return Complex(res_re, res_im)
+end
+Base.:*(z::Complex{Bool}, x::TracedRNumber{T}) where {T<:Real} = x * z
+
 # Based on https://github.com/JuliaLang/julia/blob/39255d47db7657950ff1c82137ecec5a70bae622/base/float.jl#L608-L617
 function Base.mod(
     @nospecialize(x::Reactant.TracedRNumber{T}), @nospecialize(y::Reactant.TracedRNumber{T})

test/complex.jl

@@ -60,6 +60,20 @@ end
     @test isapprox(f(y), ComplexF32(2.0 - 1.0im))
 end
 
+@testset "multiplication by im" begin
+    x = ConcreteRNumber(42.0f0)
+
+    mul_im(x) = x * im
+    mul_re(x) = x * Complex(true, false)
+    mul_re_im(x) = x * Complex(true, true)
+    mul_nothing(x) = x * Complex(false, false)
+
+    @test (@jit mul_im(x)) == Complex(0.0f0, 42.0f0)
+    @test (@jit mul_re(x)) == Complex(42.0f0, 0.0f0)
+    @test (@jit mul_re_im(x)) == Complex(42.0f0, 42.0f0)
+    @test (@jit mul_nothing(x)) == Complex(0.0f0, 0.0f0)
+end
+
 @testset "complex reduction" begin
     x = Reactant.TestUtils.construct_test_array(ComplexF32, 10, 10)
     x_ra = Reactant.to_rarray(x)