test: more fixes

Author: avik-pal

test/basic.jl

@@ -824,16 +824,19 @@ end
 end
 
 @testset "signbit" begin
-    for x in (-4, -3.14, -0.0f0, 0.0, 0, 5, 6.28f0)
-        @test @jit(signbit(ConcreteRNumber(x))) == signbit(x) broken = RunningOnTPU
+    @testset "$(typeof(x))" for x in (-4, -3.14, -0.0f0, 0.0, 0, 5, 6.28f0)
+        @test @jit(signbit(ConcreteRNumber(x))) == signbit(x) broken =
+            RunningOnTPU && eltype(x) == Float64
     end
 end
 
 @testset "copysign" begin
-    for a in (-3.14, -2, 0.0, 2.71, 42), b in (-7, -0.57, -0.0, 1, 3.14)
+    @testset "$(typeof(a)) $(typeof(b))" for a in (-3.14, -2, 0.0, 2.71, 42),
+        b in (-7, -0.57, -0.0, 1, 3.14)
         # Make sure also the return type is correct
         @test Reactant.to_number(@jit(copysign(ConcreteRNumber(a), ConcreteRNumber(b)))) ===
-            copysign(a, b) broken = RunningOnTPU
+            copysign(a, b) broken =
+            RunningOnTPU && (eltype(a) == Float64 || eltype(b) == Float64)
     end
 end
 
@@ -930,10 +933,10 @@ end
 end
 
 @testset "@code_xla" begin
-    x_ra = Reactant.to_rarray(ones(4))
+    x_ra = Reactant.to_rarray(ones(Float32, 4))
     hlo = repr(@code_xla(sin.(x_ra)))
-    @test contains(hlo, "HloModule") broken = RunningOnTPU
-    @test contains(hlo, "sine") broken = RunningOnTPU
+    @test contains(hlo, "HloModule")
+    @test contains(hlo, "sine")
 end
 
 @testset "Raise keyword" begin
@@ -1125,7 +1128,7 @@ end
         Reactant.MLIR.IR.DUMP_MLIR_DIR[] = dir
         @compile sin.(Reactant.to_rarray(Float32[1.0]))
         for mod in readdir(dir; join=true)
-            @test contains(read(mod, String), "hlo.sine") broken = RunningOnTPU
+            @test contains(read(mod, String), "hlo.sine")
         end
     end
 
@@ -1135,7 +1138,7 @@ end
         @compile exp.(Reactant.to_rarray(Float32[1.0]))
         # Make sure we don't save anything to file when compilation is
         # successful and `DUMP_MLIR_ALWAYS=false`.
-        @test isempty(readdir(dir; join=true)) broken = RunningOnTPU
+        @test isempty(readdir(dir; join=true))
     end
 
     Reactant.MLIR.IR.DUMP_MLIR_ALWAYS[] = always_old
@@ -1267,7 +1270,7 @@ accum_fn(x, y) = abs2(x) + abs2(y)
 
     @testset "accumulate" begin
         @test @jit(accumulate(accum_fn, a_ra; init=0.0f0)) ≈
-            accumulate(accum_fn, a; init=0.0f0)
+            accumulate(accum_fn, a; init=0.0f0) broken = RunningOnTPU
 
         @test @jit(accumulate(accum_fn, b_ra; init=0.0f0, dims=1)) ≈
             accumulate(accum_fn, b; dims=1, init=0.0f0) broken = RunningOnTPU

test/ops.jl

@@ -13,11 +13,11 @@ const RunningOnAppleX86 = Sys.isapple() && Sys.ARCH === :x86_64
     x = Reactant.to_rarray([1.0, -1.0])
     @test [1.0, 1.0] ≈ @jit Ops.abs(x)
 
-    x = Reactant.to_rarray([
-        3.0+4im -3.0+4im
-        3.0-4im -3.0-4im
-    ])
     @test begin
+        x = Reactant.to_rarray([
+            3.0+4im -3.0+4im
+            3.0-4im -3.0-4im
+        ])
         [
             5.0 5.0
             5.0 5.0
@@ -38,11 +38,13 @@ end
     b = Reactant.to_rarray([5.5, 6.6, -7.7, -8.8])
     @test Array(a) .+ Array(b) ≈ @jit Ops.add(a, b)
 
-    a = Reactant.to_rarray([1.1 + 2.2im, 3.3 + 4.4im, 5.5 + 6.6im, 7.7 + 8.8im])
-    b = Reactant.to_rarray([
-        9.9 + 10.10im, 11.11 + 12.12im, -13.13 + -14.14im, -15.15 + -16.16im
-    ])
-    @test Array(a) .+ Array(b) ≈ @jit(Ops.add(a, b)) broken = RunningOnTPU
+    @test begin
+        a = Reactant.to_rarray([1.1 + 2.2im, 3.3 + 4.4im, 5.5 + 6.6im, 7.7 + 8.8im])
+        b = Reactant.to_rarray([
+            9.9 + 10.10im, 11.11 + 12.12im, -13.13 + -14.14im, -15.15 + -16.16im
+        ])
+        Array(a) .+ Array(b) ≈ @jit(Ops.add(a, b))
+    end broken = RunningOnTPU
 end
 
 @testset "after_all" begin
@@ -100,16 +102,16 @@ end
     @test cholesky(Array(x)).U ≈ @jit g1(x)
     @test transpose(cholesky(Array(x)).U) ≈ @jit g2(x)
 
-    x = Reactant.to_rarray(
-        [
-            10.0+0.0im 2.0-3.0im 3.0-4.0im
-            2.0+3.0im 5.0+0.0im 3.0-2.0im
-            3.0+4.0im 3.0+2.0im 9.0+0.0im
-        ],
-    )
-
-    @test cholesky(Array(x)).U ≈ @jit(g1(x)) broken = RunningOnTPU
-    @test adjoint(cholesky(Array(x)).U) ≈ @jit(g2(x)) broken = RunningOnTPU
+    @test begin
+        x = Reactant.to_rarray(
+            [
+                10.0+0.0im 2.0-3.0im 3.0-4.0im
+                2.0+3.0im 5.0+0.0im 3.0-2.0im
+                3.0+4.0im 3.0+2.0im 9.0+0.0im
+            ],
+        )
+        cholesky(Array(x)).U ≈ @jit(g1(x)) && adjoint(cholesky(Array(x)).U) ≈ @jit(g2(x))
+    end broken = RunningOnTPU
 end
 
 @testset "clamp" begin

test/runtests.jl

@@ -13,29 +13,18 @@ const REACTANT_TEST_GROUP = lowercase(get(ENV, "REACTANT_TEST_GROUP", "all"))
         end
 
         # @safetestset "Layout" include("layout.jl")
-        # @info "Layout tests finished"
         # @safetestset "Tracing" include("tracing.jl")
-        # @info "Tracing tests finished"
         @safetestset "Basic" include("basic.jl")
         @info "Basic tests finished"
         # @safetestset "Constructor" include("constructor.jl")
-        # @info "Constructor tests finished"
-        @safetestset "Autodiff" include("autodiff.jl")
-        @info "Autodiff tests finished"
-        @safetestset "Complex" include("complex.jl")
-        @info "Complex tests finished"
-        @safetestset "Broadcast" include("bcast.jl")
-        @info "Broadcast tests finished"
-        @safetestset "Struct" include("struct.jl")
-        @info "Struct tests finished"
-        @safetestset "Closure" include("closure.jl")
-        @info "Closure tests finished"
-        @safetestset "Compile" include("compile.jl")
-        @info "Compile tests finished"
-        @safetestset "IR" include("ir.jl")
-        @info "IR tests finished"
-        @safetestset "Buffer Donation" include("buffer_donation.jl")
-        @info "Buffer Donation tests finished"
+        # @safetestset "Autodiff" include("autodiff.jl")
+        # @safetestset "Complex" include("complex.jl")
+        # @safetestset "Broadcast" include("bcast.jl")
+        # @safetestset "Struct" include("struct.jl")
+        # @safetestset "Closure" include("closure.jl")
+        # @safetestset "Compile" include("compile.jl")
+        # @safetestset "IR" include("ir.jl")
+        # @safetestset "Buffer Donation" include("buffer_donation.jl")
         @safetestset "Shortcuts to MLIR ops" include("ops.jl")
         @info "Shortcuts to MLIR ops tests finished"
         @safetestset "Wrapped Arrays" include("wrapped_arrays.jl")