Fix irrational promotion (#924)

* Fix irrational promotion

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* fixup

* fix

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update basic.jl

* fix

* fix test

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: wsmoses

src/TracedRNumber.jl

@@ -16,6 +16,10 @@ function Base.eps(::Type{TracedRNumber{T}}) where {T}
     return TracedUtils.promote_to(TracedRNumber{T}, eps(T))
 end
 
+function Base.rtoldefault(T::Type{<:TracedRNumber})
+    return T(Base.rtoldefault(unwrapped_eltype(T)))
+end
+
 function Base.isfinite(x::TracedRNumber{<:Complex})
     return isfinite(real(x)) & isfinite(imag(x))
 end
@@ -59,6 +63,18 @@ function Base.promote_rule(::Type{TracedRNumber{T}}, ::Type{S}) where {T,S}
     return TracedRNumber{Base.promote_type(T, S)}
 end
 
+function Base.promote_rule(
+    T::Type{<:AbstractIrrational}, ::Type{Reactant.TracedRNumber{S}}
+) where {S}
+    return TracedRNumber{Base.promote_type(T, S)}
+end
+
+function Base.promote_rule(
+    ::Type{Reactant.TracedRNumber{S}}, T::Type{<:AbstractIrrational}
+) where {S}
+    return TracedRNumber{Base.promote_type(T, S)}
+end
+
 # NOTE: This is inconsistent with the behavior of `convert` but we do it since it is a very
 #       common usecase
 TracedRNumber{T}(x::TracedRNumber{T}) where {T} = x

src/TracedUtils.jl

@@ -651,7 +651,13 @@ function broadcast_to_size(arg::Base.RefValue, rsize)
     return arg
 end
 
-broadcast_to_size(arg::Number, rsize) = Ops.constant(Base.fill(arg, Tuple(rsize)))
+function broadcast_to_size(arg::AbstractIrrational, rsize)
+    return broadcast_to_size(Base.convert(Float64, arg), rsize)
+end
+
+function broadcast_to_size(arg::ReactantPrimitive, rsize)
+    return Ops.constant(Base.fill(arg, Tuple(rsize)))
+end
 
 function broadcast_to_size(arg::TracedRNumber{T}, rsize) where {T}
     length(rsize) == 0 && return arg

test/basic.jl

@@ -998,3 +998,11 @@ end
     @test res[2] == 215
     @test res[3] == 216
 end
+
+mulpi(x) = π * x
+
+@testset "Irrational promotion" begin
+    x = Reactant.to_rarray(ones(2))
+    y = @jit mulpi(x)
+    @test all(Array(y) .≈ π)
+end