test: mark tests as broken instead of skipping them

Author: avik-pal

test/autodiff.jl

@@ -195,20 +195,18 @@ end
     @test res2 ≈ 4 * 3 * 3.1^2
 end
 
-if !contains(string(Reactant.devices()[1]), "TPU")
-    @testset "Seed initialization of Complex arrays on matmul: Issue #593" begin
-        a = ones(ComplexF64, 2, 2)
-        b = 2.0 * ones(ComplexF64, 2, 2)
-        a_re = Reactant.to_rarray(a)
-        b_re = Reactant.to_rarray(b)
-        df(x, y) = Enzyme.gradient(ReverseWithPrimal, *, x, y)
-        @test begin
-            res = @jit df(a_re, b_re) # before, this segfaulted
-            (res.val ≈ 4ones(2, 2)) &&
-                (res.derivs[1] ≈ 4ones(2, 2)) &&
-                (res.derivs[2] ≈ 2ones(2, 2))
-        end
-    end
+@testset "Seed initialization of Complex arrays on matmul: Issue #593" begin
+    a = ones(ComplexF64, 2, 2)
+    b = 2.0 * ones(ComplexF64, 2, 2)
+    a_re = Reactant.to_rarray(a)
+    b_re = Reactant.to_rarray(b)
+    df(x, y) = Enzyme.gradient(ReverseWithPrimal, *, x, y)
+    @test begin
+        res = @jit df(a_re, b_re) # before, this segfaulted
+        (res.val ≈ 4ones(2, 2)) &&
+            (res.derivs[1] ≈ 4ones(2, 2)) &&
+            (res.derivs[2] ≈ 2ones(2, 2))
+    end broken=contains(string(Reactant.devices()[1]), "TPU")
 end
 
 @testset "onehot" begin

test/basic.jl

@@ -1,13 +1,9 @@
-using Reactant
-using Test
-using Enzyme
-using Statistics
-using Random
+using Reactant, Test, Enzyme, Statistics, Random, InteractiveUtils
 Random.seed!(123)
 
-fastmax(x::AbstractArray{T}) where {T} = reduce(max, x; dims=1, init=float(T)(-Inf))
+const RunningOnTPU = contains(string(Reactant.devices()[1]), "TPU")
 
-using InteractiveUtils
+fastmax(x::AbstractArray{T}) where {T} = reduce(max, x; dims=1, init=float(T)(-Inf))
 
 @testset "2D sum" begin
     x = rand(2, 10)
@@ -420,12 +416,12 @@ end
 
 @testset "Complex runtime: $CT" for CT in (ComplexF32, ComplexF64)
     # complex f64 not supported on tpu
-    if CT == ComplexF32 || !contains(string(Reactant.devices()[1]), "TPU")
-        a = Reactant.to_rarray(ones(CT, 2))
-        b = Reactant.to_rarray(ones(CT, 2))
+    a = Reactant.to_rarray(ones(CT, 2))
+    b = Reactant.to_rarray(ones(CT, 2))
+    @test begin
         c = Reactant.compile(+, (a, b))(a, b)
-        @test c == ones(CT, 2) + ones(CT, 2)
-    end
+        c == ones(CT, 2) + ones(CT, 2)
+    end broken = CT != ComplexF32 && RunningOnTPU
 end
 
 @testset "Scalars" begin
@@ -784,20 +780,20 @@ end
     x = Reactant.to_rarray([1.0, NaN, Inf, -Inf, NaN])
     @test @jit(isfinite.(x)) == [true, false, false, false, false]
 
-    if !contains(string(Reactant.devices()[1]), "TPU")
+    @test begin
         x = Reactant.to_rarray([1.0, NaN, Inf, -Inf, NaN] .* im)
-        @test @jit(isfinite.(x)) == [true, false, false, false, false]
-    end
+        @jit(isfinite.(x)) == [true, false, false, false, false]
+    end broken = RunningOnTPU
 end
 
 @testset "isnan" begin
     x = Reactant.to_rarray([1.0, NaN, Inf, -Inf, NaN])
     @test @jit(isnan.(x)) == [false, true, false, false, true]
 
-    if !contains(string(Reactant.devices()[1]), "TPU")
+    @test begin
         x = Reactant.to_rarray([1.0, NaN, Inf, -Inf, NaN] .* im)
-        @test @jit(isnan.(x)) == [false, true, false, false, true]
-    end
+        @jit(isnan.(x)) == [false, true, false, false, true]
+    end broken = RunningOnTPU
 end
 
 @testset "isnan/isfinite" begin
@@ -820,11 +816,10 @@ end
     b = [6.6, -2.2, -8.8, 4.4, -10.1]
 
     expected_mod = mod.(a, b)
-    if !contains(string(Reactant.devices()[1]), "TPU")
-        @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
-    end
+    @test @jit(mod.(Reactant.to_rarray(a), Reactant.to_rarray(b))) ≈ expected_mod broken =
+        RunningOnTPU
+    @test @jit(mod.(a, Reactant.to_rarray(b))) ≈ expected_mod broken = RunningOnTPU
+    @test @jit(mod.(Reactant.to_rarray(a), b)) ≈ expected_mod broken = RunningOnTPU
 
     expected_rem = rem.(a, b)
     @test @jit(rem.(Reactant.to_rarray(a), Reactant.to_rarray(b))) ≈ expected_rem
@@ -838,22 +833,17 @@ end
     end
 end
 
-if !contains(string(Reactant.devices()[1]), "TPU")
-    @testset "signbit" begin
-        for x in (-4, -3.14, -0.0f0, 0.0, 0, 5, 6.28f0)
-            @test @jit(signbit(ConcreteRNumber(x))) == signbit(x)
-        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
     end
 end
 
-if !contains(string(Reactant.devices()[1]), "TPU")
-    @testset "copysign" begin
-        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)
-        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)
+        # Make sure also the return type is correct
+        @test Reactant.to_number(@jit(copysign(ConcreteRNumber(a), ConcreteRNumber(b)))) ===
+            copysign(a, b) broken = RunningOnTPU
     end
 end
 
@@ -949,13 +939,11 @@ end
     ra[:a] ≈ (2.7 * 2) * ones(4)
 end
 
-if !contains(string(Reactant.devices()[1]), "TPU")
-    @testset "@code_xla" begin
-        x_ra = Reactant.to_rarray(ones(4))
-        hlo = repr(@code_xla(sin.(x_ra)))
-        @test contains(hlo, "HloModule")
-        @test contains(hlo, "sine")
-    end
+@testset "@code_xla" begin
+    x_ra = Reactant.to_rarray(ones(4))
+    hlo = repr(@code_xla(sin.(x_ra)))
+    @test contains(hlo, "HloModule") broken = RunningOnTPU
+    @test contains(hlo, "sine") broken = RunningOnTPU
 end
 
 @testset "Raise keyword" begin
@@ -999,14 +987,11 @@ end
     @test Array(x) ≈ Array(y) ./ 2
 end
 
-if !contains(string(Reactant.devices()[1]), "TPU")
-    @testset "Hlo Cost Analysis" begin
-        x_ra = Reactant.to_rarray(rand(4, 4))
-        mul_comp = @compile x_ra * x_ra
-        cost = Reactant.XLA.cost_analysis(mul_comp)
-
-        @test cost isa Reactant.XLA.HloCostAnalysisProperties
-    end
+@test "HLO Cost Analysis" begin
+    x_ra = Reactant.to_rarray(rand(4, 4))
+    mul_comp = @compile x_ra * x_ra
+    @test Reactant.XLA.cost_analysis(mul_comp) isa Reactant.XLA.HloCostAnalysisProperties broken =
+        RunningOnTPU
 end
 
 function fractional_idx(times, t)
@@ -1140,32 +1125,30 @@ end
     end
 end
 
-if !contains(string(Reactant.devices()[1]), "TPU")
-    @testset "Dump MLIR modules" begin
-        always_old = Reactant.MLIR.IR.DUMP_MLIR_ALWAYS[]
-        dir_old = Reactant.MLIR.IR.DUMP_MLIR_DIR[]
+@testset "Dump MLIR modules" begin
+    always_old = Reactant.MLIR.IR.DUMP_MLIR_ALWAYS[]
+    dir_old = Reactant.MLIR.IR.DUMP_MLIR_DIR[]
 
-        mktempdir() do dir
-            Reactant.MLIR.IR.DUMP_MLIR_ALWAYS[] = true
-            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")
-            end
-        end
-
-        mktempdir() do dir
-            Reactant.MLIR.IR.DUMP_MLIR_ALWAYS[] = false
-            Reactant.MLIR.IR.DUMP_MLIR_DIR[] = dir
-            @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))
+    mktempdir() do dir
+        Reactant.MLIR.IR.DUMP_MLIR_ALWAYS[] = true
+        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
         end
+    end
 
-        Reactant.MLIR.IR.DUMP_MLIR_ALWAYS[] = always_old
-        Reactant.MLIR.IR.DUMP_MLIR_DIR[] = dir_old
+    mktempdir() do dir
+        Reactant.MLIR.IR.DUMP_MLIR_ALWAYS[] = false
+        Reactant.MLIR.IR.DUMP_MLIR_DIR[] = dir
+        @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
     end
+
+    Reactant.MLIR.IR.DUMP_MLIR_ALWAYS[] = always_old
+    Reactant.MLIR.IR.DUMP_MLIR_DIR[] = dir_old
 end
 
 @testset "Allocator Stats" begin