fix: unwanted promotion in copysign (#1590)

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.158"
+version = "0.2.159"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/TracedRNumber.jl

@@ -171,15 +171,27 @@ function Base.mod(
     r = rem(x, y)
     return ifelse(r == 0, copysign(r, y), ifelse((r > 0) ⊻ (y > 0), r + y, r))
 end
-function Base.mod(
-    @nospecialize(lhs::TracedRNumber{T}), @nospecialize(rhs::Number)
+
+function Base.mod1(
+    @nospecialize(x::Reactant.TracedRNumber{T}), @nospecialize(y::Reactant.TracedRNumber{T})
 ) where {T}
-    return mod(lhs, TracedUtils.promote_to(TracedRNumber{T}, rhs))
+    m = mod(x, y)
+    return ifelse(m == 0, y, m)
 end
-function Base.mod(
-    @nospecialize(lhs::Number), @nospecialize(rhs::TracedRNumber{T})
-) where {T}
-    return mod(TracedUtils.promote_to(TracedRNumber{T}, lhs), rhs)
+
+for op in (:mod, :mod1)
+    @eval begin
+        function Base.$op(
+            @nospecialize(lhs::TracedRNumber{T}), @nospecialize(rhs::Number)
+        ) where {T}
+            return mod(lhs, TracedUtils.promote_to(TracedRNumber{T}, rhs))
+        end
+        function Base.$op(
+            @nospecialize(lhs::Number), @nospecialize(rhs::TracedRNumber{T})
+        ) where {T}
+            return mod(TracedUtils.promote_to(TracedRNumber{T}, lhs), rhs)
+        end
+    end
 end
 
 function Base.div(@nospecialize(lhs::TracedRNumber{T}), rhs) where {T<:Integer}
@@ -938,7 +950,9 @@ for (Ti, Tf) in ((Int16, Float16), (Int32, Float32), (Int64, Float64))
 end
 Base.signbit(::TracedRNumber{<:Unsigned}) = ConcretePJRTNumber(false)
 
-Base.copysign(x::TracedRNumber, y::TracedRNumber) = ifelse(signbit(y), -1, 1) * abs(x)
+function Base.copysign(x::TracedRNumber, y::TracedRNumber)
+    return ifelse(signbit(y), -one(x), one(x)) * abs(x)
+end
 function Base.copysign(x::TracedRNumber{T}, y::S) where {T,S<:Number}
     return copysign(x, TracedUtils.promote_to(TracedRNumber{S}, y))
 end

test/basic.jl

@@ -1514,3 +1514,21 @@ stack_numbers(x) = stack([sum(x[:, i]) for i in axes(x, 2)])
 
     @test @jit(stack_numbers(x_ra)) ≈ stack_numbers(x)
 end
+
+@testset "copysign/mod type check" begin
+    x = ConcreteRNumber(Int32(5))
+    y = ConcreteRNumber(Int32(3))
+    @test @jit(copysign(x, y)) isa ConcreteRNumber{Int32}
+    @test @jit(mod(x, y)) isa ConcreteRNumber{Int32}
+end
+
+@testset "mod1" begin
+    x = collect(Int32, 1:12)
+    y = Int32(10)
+
+    @testset for xᵢ in x
+        res = @jit mod1(ConcreteRNumber(xᵢ), ConcreteRNumber(y))
+        @test res isa ConcreteRNumber{Int32}
+        @test res == mod1(xᵢ, y)
+    end
+end