test: mark tests as broken instead of skipping them (#1536)

* test: mark tests as broken/skip

* test: more fixes

* test: more fixes

* test: more fixes

* test: nn tests

* ci: only run failing cases (drop me)

* test: print

* test: reduce test configs

* test: try adjusting precision

* ci: restore runners

* ci: run under gdb (drop me)

* ci: remove gdb

* fix: temporarily remove conditional in softmax

* test: with highest precision for accuracy

* chore: fmt

Author: avik-pal

.github/workflows/CI.yml

@@ -66,6 +66,16 @@ jobs:
             assertions: false
             test_group: core
             runtime: "IFRT"
+          - os: linux-x86-ct6e-180-4tpu
+            version: "1.11"
+            assertions: false
+            test_group: integration
+            runtime: "IFRT"
+          - os: linux-x86-ct6e-180-4tpu
+            version: "1.11"
+            assertions: false
+            test_group: neural_networks
+            runtime: "IFRT"
           - os: ubuntu-24.04
             version: "1.10"
             assertions: true
@@ -81,18 +91,18 @@ jobs:
             assertions: true
             test_group: integration
             runtime: "PJRT"
-          # - os: ubuntu-24.04
-          #   libReactant: packaged
-          #   version: '1.10'
-          #   test_group: core
-          # - os: ubuntu-24.04
-          #   libReactant: packaged
-          #   version: '1.10'
-          #   test_group: neural_networks
-          # - os: ubuntu-24.04
-          #   libReactant: packaged
-          #   version: '1.10'
-          #   test_group: integration
+          - os: ubuntu-24.04
+            libReactant: packaged
+            version: '1.10'
+            test_group: core
+          - os: ubuntu-24.04
+            libReactant: packaged
+            version: '1.10'
+            test_group: neural_networks
+          - os: ubuntu-24.04
+            libReactant: packaged
+            version: '1.10'
+            test_group: integration
     steps:
       - name: Set TMPDIR
         # We have to use `${GITHUB_WORKSPACE}` instead of `github.workspace` because GitHub

ext/ReactantNNlibExt/Implementations.jl

@@ -10,11 +10,13 @@ function NNlib.softmax!(out::AnyTracedRArray{T,N}, x::AbstractArray; dims=1) whe
     x = T.(Reactant.materialize_traced_array(x))
     max_ = maximum(x; dims)
     diff = exp.(x .- max_)
-    @trace if all(isfinite, max_)
-        @. out = diff
-    else
-        @. out = ifelse(isinf(max_), ifelse(isinf(x), T(1), T(0)), diff)
-    end
+    # TOOD: re-enable conditional once https://github.com/EnzymeAD/Reactant.jl/issues/1581
+    # fixed
+    # @trace if all(isfinite, max_)
+    @. out = diff
+    # else
+    #     @. out = ifelse(isinf(max_), ifelse(isinf(x), T(1), T(0)), diff)
+    # end
     out ./= sum(out; dims)
     return out
 end
@@ -23,11 +25,13 @@ function NNlib.logsoftmax!(out::AnyTracedRArray{T}, x::AbstractArray; dims=1) wh
     x = T.(Reactant.materialize_traced_array(x))
     max_ = maximum(x; dims)
     diff = x .- max_
-    @trace if all(isfinite, max_)
-        @. out = diff
-    else
-        @. out = ifelse(isinf(max_), ifelse(isinf(x), T(0), -T(Inf)), diff)
-    end
+    # TOOD: re-enable conditional once https://github.com/EnzymeAD/Reactant.jl/issues/1581
+    # fixed
+    # @trace if all(isfinite, max_)
+    @. out = diff
+    # else
+    #     @. out = ifelse(isinf(max_), ifelse(isinf(x), T(0), -T(Inf)), diff)
+    # end
     out .-= log.(sum(exp, out; dims))
     return out
 end

src/accelerators/TPU.jl

@@ -43,7 +43,7 @@ function download_libtpu_if_needed(path=nothing)
         zip_file_path = joinpath(path, "tpu.zip")
         tmp_dir = joinpath(path, "tmp")
         Downloads.download(
-            "https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu-nightly/libtpu_nightly-0.1.dev20250727+nightly-py3-none-manylinux_2_31_x86_64.whl",
+            "https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu-nightly/libtpu_nightly-0.1.dev20250811+nightly-py3-none-manylinux_2_31_x86_64.whl",
             zip_file_path,
         )
         run(`$(unzip()) -qq $(zip_file_path) -d $(tmp_dir)`)

src/xla/XLA.jl

@@ -131,7 +131,8 @@ function __init__()
             XLA_REACTANT_GPU_MEM_FRACTION[] = parse(
                 Float64, ENV["XLA_REACTANT_GPU_MEM_FRACTION"]
             )
-            @debug "XLA_REACTANT_GPU_MEM_FRACTION: " XLA_REACTANT_GPU_MEM_FRACTION[]
+            @debug "XLA_REACTANT_GPU_MEM_FRACTION: " XLA_REACTANT_GPU_MEM_FRACTION[] maxlog =
+                1
             if XLA_REACTANT_GPU_MEM_FRACTION[] > 1 || XLA_REACTANT_GPU_MEM_FRACTION[] < 0
                 error("XLA_REACTANT_GPU_MEM_FRACTION must be between 0 and 1")
             end
@@ -141,16 +142,18 @@ function __init__()
             XLA_REACTANT_GPU_PREALLOCATE[] = parse(
                 Bool, ENV["XLA_REACTANT_GPU_PREALLOCATE"]
             )
-            @debug "XLA_REACTANT_GPU_PREALLOCATE: " XLA_REACTANT_GPU_PREALLOCATE[]
+            @debug "XLA_REACTANT_GPU_PREALLOCATE: " XLA_REACTANT_GPU_PREALLOCATE[] maxlog =
+                1
         end
 
         if haskey(ENV, "REACTANT_VISIBLE_GPU_DEVICES")
             global_state.local_gpu_device_ids =
                 parse.(Int, split(ENV["REACTANT_VISIBLE_GPU_DEVICES"], ","))
-            @debug "REACTANT_VISIBLE_GPU_DEVICES: " global_state.local_gpu_device_ids
+            @debug "REACTANT_VISIBLE_GPU_DEVICES: " global_state.local_gpu_device_ids maxlog =
+                1
         end
 
-        @debug "REACTANT_XLA_RUNTIME: " REACTANT_XLA_RUNTIME
+        @debug "REACTANT_XLA_RUNTIME: " REACTANT_XLA_RUNTIME maxlog = 1
 
         @ccall MLIR.API.mlir_c.RegisterEnzymeXLACPUHandler()::Cvoid
         @ccall MLIR.API.mlir_c.RegisterEnzymeXLAGPUHandler()::Cvoid

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
+@testset "Seed initialization of Complex arrays on matmul: Issue #593" begin
+    df(x, y) = Enzyme.gradient(ReverseWithPrimal, *, x, y)
+    @test 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
+        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 skip = contains(string(Reactant.devices()[1]), "TPU")
 end
 
 @testset "onehot" begin
@@ -257,7 +255,7 @@ end
 
 @testset "seed" begin
     x = Reactant.to_rarray(rand(2, 2))
-    st = (; rng=Reactant.ConcreteRNG())
+    st = (; rng=Reactant.ReactantRNG())
 
     @test begin
         hlo = @code_hlo gradient_fn(x, st)