feat: add isinf dispatches (#826)

Author: avik-pal

src/Ops.jl

@@ -280,20 +280,19 @@ for (dialect, op) in [
 end
 
 # is* checks
-for (dialect, op) in [
-    #(:stablehlo, :is_finite),
-    (:chlo, :is_inf),
-    (:chlo, :is_neg_inf),
-    (:chlo, :is_pos_inf),
-]
+for (dialect, op) in
+    [(:stablehlo, :is_finite), (:chlo, :is_inf), (:chlo, :is_neg_inf), (:chlo, :is_pos_inf)]
+    result = dialect == :stablehlo ? :y : :result
     @eval begin
         @noinline function $op(
             x::TracedRArray{T,N};
             location=mlir_stacktrace($(string(op)), @__FILE__, @__LINE__),
         ) where {T,N}
             res = MLIR.IR.result(
                 $(:($dialect.$op))(
-                    x.mlir_data; result=mlir_type(TracedRArray{Bool,N}, size(x)), location
+                    x.mlir_data;
+                    $(result)=mlir_type(TracedRArray{Bool,N}, size(x)),
+                    location,
                 ),
             )
             return TracedRArray{Bool,N}((), res, size(x))
@@ -305,34 +304,14 @@ for (dialect, op) in [
         ) where {T}
             res = MLIR.IR.result(
                 $(:($dialect.$op))(
-                    x.mlir_data; result=mlir_type(TracedRArray{Bool,0}, ()), location
+                    x.mlir_data; $(result)=mlir_type(TracedRArray{Bool,0}, ()), location
                 ),
             )
             return TracedRNumber{Bool}((), res)
         end
     end
 end
 
-@noinline function is_finite(
-    x::TracedRArray{T,N}; location=mlir_stacktrace("is_finite", @__FILE__, @__LINE__)
-) where {T,N}
-    res = MLIR.IR.result(
-        stablehlo.is_finite(
-            x.mlir_data; y=mlir_type(TracedRArray{Bool,N}, size(x)), location
-        ),
-    )
-    return TracedRArray{Bool,N}((), res, size(x))
-end
-
-@noinline function is_finite(
-    x::TracedRNumber{T}; location=mlir_stacktrace("is_finite", @__FILE__, @__LINE__)
-) where {T}
-    res = MLIR.IR.result(
-        stablehlo.is_finite(x.mlir_data; y=mlir_type(TracedRArray{Bool,0}, ()), location)
-    )
-    return TracedRNumber{Bool}((), res)
-end
-
 # fixes to default automated implementations
 @noinline function abs(
     x::TracedRArray{Complex{T},N}; location=mlir_stacktrace("abs", @__FILE__, @__LINE__)

src/TracedRNumber.jl

@@ -19,16 +19,18 @@ end
 function Base.isfinite(x::TracedRNumber{<:Complex})
     return isfinite(real(x)) & isfinite(imag(x))
 end
-function Base.isfinite(x::TracedRNumber{T}) where {T<:AbstractFloat}
-    return Reactant.Ops.is_finite(x)
-end
+Base.isfinite(x::TracedRNumber{<:AbstractFloat}) = Ops.is_finite(x)
 
-function Base.isnan(x::TracedRNumber{T}) where {T<:AbstractFloat}
-    return !isfinite(x) & (x != typemax(T)) & (x != typemin(T))
-end
 function Base.isnan(x::TracedRNumber{<:Complex})
     return isnan(real(x)) | isnan(imag(x))
 end
+function Base.isnan(x::TracedRNumber{T}) where {T<:AbstractFloat}
+    return !isfinite(x) & (x != typemax(T)) & (x != typemin(T))
+end
+
+Base.isinf(x::TracedRNumber{<:Complex}) = isinf(real(x)) | isinf(imag(x))
+Base.isinf(x::TracedRNumber{<:AbstractFloat}) = Ops.is_inf(x)
+Base.isinf(::TracedRNumber{<:Integer}) = false
 
 function Base.show(io::IOty, X::TracedRNumber{T}) where {T,IOty<:Union{IO,IOContext}}
     return print(io, "TracedRNumber{", T, "}(", X.paths, ")")

test/basic.jl

@@ -766,18 +766,18 @@ end
 
 @testset "isfinite" begin
     x = Reactant.to_rarray([1.0, NaN, Inf, -Inf, NaN])
-    @test Reactant.@jit(isfinite.(x)) == [true, false, false, false, false]
+    @test @jit(isfinite.(x)) == [true, false, false, false, false]
 
     x = Reactant.to_rarray([1.0, NaN, Inf, -Inf, NaN] .* im)
-    @test Reactant.@jit(isfinite.(x)) == [true, false, false, false, false]
+    @test @jit(isfinite.(x)) == [true, false, false, false, false]
 end
 
 @testset "isnan" begin
     x = Reactant.to_rarray([1.0, NaN, Inf, -Inf, NaN])
-    @test Reactant.@jit(isnan.(x)) == [false, true, false, false, true]
+    @test @jit(isnan.(x)) == [false, true, false, false, true]
 
     x = Reactant.to_rarray([1.0, NaN, Inf, -Inf, NaN] .* im)
-    @test Reactant.@jit(isnan.(x)) == [false, true, false, false, true]
+    @test @jit(isnan.(x)) == [false, true, false, false, true]
 end
 
 @testset "isnan/isfinite" begin
@@ -787,19 +787,27 @@ end
     @test !isfinite(Reactant.to_rarray(Inf; track_numbers=Number))
 end
 
+@testset "isinf" begin
+    @test Bool(@jit(isinf(ConcreteRNumber(Inf))))
+    @test Bool(@jit(isinf(ConcreteRNumber(-Inf))))
+    @test !Bool(@jit(isinf(ConcreteRNumber(2))))
+    @test !Bool(@jit(isinf(ConcreteRNumber(2.0))))
+    @test !Bool(@jit(isinf(ConcreteRNumber(true))))
+end
+
 @testset "mod and rem" begin
     a = [-1.1, 7.7, -3.3, 9.9, -5.5]
     b = [6.6, -2.2, -8.8, 4.4, -10.1]
 
     expected_mod = mod.(a, b)
-    @test Reactant.@jit(mod.(Reactant.to_rarray(a), Reactant.to_rarray(b))) ≈ expected_mod
-    @test Reactant.@jit(mod.(a, Reactant.to_rarray(b))) ≈ expected_mod
-    @test Reactant.@jit(mod.(Reactant.to_rarray(a), b)) ≈ expected_mod
+    @test @jit(mod.(Reactant.to_rarray(a), Reactant.to_rarray(b))) ≈ expected_mod
+    @test @jit(mod.(a, Reactant.to_rarray(b))) ≈ expected_mod
+    @test @jit(mod.(Reactant.to_rarray(a), b)) ≈ expected_mod
 
     expected_rem = rem.(a, b)
-    @test Reactant.@jit(rem.(Reactant.to_rarray(a), Reactant.to_rarray(b))) ≈ expected_rem
-    @test Reactant.@jit(rem.(a, Reactant.to_rarray(b))) ≈ expected_rem
-    @test Reactant.@jit(rem.(Reactant.to_rarray(a), b)) ≈ expected_rem
+    @test @jit(rem.(Reactant.to_rarray(a), Reactant.to_rarray(b))) ≈ expected_rem
+    @test @jit(rem.(a, Reactant.to_rarray(b))) ≈ expected_rem
+    @test @jit(rem.(Reactant.to_rarray(a), b)) ≈ expected_rem
 end
 
 @testset "xor" begin
@@ -910,7 +918,7 @@ end
     x[:b] = 3.1 * ones(4)
 
     ra = Reactant.to_rarray(x)
-    Reactant.@jit dip!(ra)
+    @jit dip!(ra)
     ra[:a] ≈ (2.7 * 2) * ones(4)
 end