Skip Reactant Core tests when Reactant_jll is unavailable (fixes x86 CI)

test/layers/layers_testsetup.jl

@@ -10,7 +10,9 @@ const LAYERS_SETUPS = [
 
 function run_op_tests(op, setups)
     rng = StableRNG(12345)
-    xdev = reactant_device(; force = true)
+    if REACTANT_AVAILABLE
+        xdev = reactant_device(; force = true)
+    end
 
     @testset "$(length(setup.m))D | shift=$(setup.shift)" for setup in setups
         in_chs = setup.x_size[end - 1]
@@ -26,20 +28,22 @@ function run_op_tests(op, setups)
         @test size(first(m(x, ps, st))) == setup.y_size
         res = first(m(x, ps, st))
 
-        ps_ra, st_ra = xdev((ps, st))
-        x_ra = xdev(x)
-        y_ra = xdev(rand(rng, Float32, setup.y_size...))
+        if REACTANT_AVAILABLE
+            ps_ra, st_ra = xdev((ps, st))
+            x_ra = xdev(x)
+            y_ra = xdev(rand(rng, Float32, setup.y_size...))
 
-        res_ra, _ = @jit m(x_ra, ps_ra, st_ra)
-        @test res_ra ≈ res atol = 1.0f-2 rtol = 1.0f-2
+            res_ra, _ = @jit m(x_ra, ps_ra, st_ra)
+            @test res_ra ≈ res atol = 1.0f-2 rtol = 1.0f-2
 
-        @testset "check gradients" begin
-            ∂x_fd, ∂ps_fd = ∇sumabs2_finite_difference(m, x, ps, st)
-            ∂x_ra, ∂ps_ra = ∇sumabs2_reactant(m, x_ra, ps_ra, st_ra)
-            ∂x_ra, ∂ps_ra = (∂x_ra, ∂ps_ra) |> cpu_device()
+            @testset "check gradients" begin
+                ∂x_fd, ∂ps_fd = ∇sumabs2_finite_difference(m, x, ps, st)
+                ∂x_ra, ∂ps_ra = ∇sumabs2_reactant(m, x_ra, ps_ra, st_ra)
+                ∂x_ra, ∂ps_ra = (∂x_ra, ∂ps_ra) |> cpu_device()
 
-            @test ∂x_fd ≈ ∂x_ra atol = 1.0f-2 rtol = 1.0f-2
-            @test check_approx(∂ps_fd, ∂ps_ra; atol = 1.0f-2, rtol = 1.0f-2)
+                @test ∂x_fd ≈ ∂x_ra atol = 1.0f-2 rtol = 1.0f-2
+                @test check_approx(∂ps_fd, ∂ps_ra; atol = 1.0f-2, rtol = 1.0f-2)
+            end
         end
     end
     return nothing

test/models/deeponet_tests.jl

@@ -24,7 +24,9 @@ include("../shared_testsetup.jl")
         ),
     ]
 
-    xdev = reactant_device(; force = true)
+    if REACTANT_AVAILABLE
+        xdev = reactant_device(; force = true)
+    end
 
     @testset "$(setup.name)" for setup in setups
         u = rand(Float32, setup.u_size...)
@@ -36,22 +38,26 @@ include("../shared_testsetup.jl")
         pred = first(deeponet((u, y), ps, st))
         @test setup.out_size == size(pred)
 
-        ps_ra, st_ra = (ps, st) |> xdev
-        u_ra, y_ra = (u, y) |> xdev
-
-        pred_ra = @jit deeponet((u_ra, y_ra), ps_ra, st_ra)
-        @test first(pred_ra) ≈ pred atol = 1.0f-2 rtol = 1.0f-2
-
-        @testset "check gradients" begin
-            (∂u_fd, ∂y_fd), ∂ps_fd = ∇sumabs2_finite_difference(
-                deeponet, (u, y), ps, st
-            )
-            (∂u_ra, ∂y_ra), ∂ps_ra = ∇sumabs2_reactant(deeponet, (u_ra, y_ra), ps_ra, st_ra)
-            (∂u_ra, ∂y_ra), ∂ps_ra = ((∂u_ra, ∂y_ra), ∂ps_ra) |> cpu_device()
-
-            @test ∂u_fd ≈ ∂u_ra atol = 1.0f-2 rtol = 1.0f-2
-            @test ∂y_fd ≈ ∂y_ra atol = 1.0f-2 rtol = 1.0f-2
-            @test check_approx(∂ps_fd, ∂ps_ra; atol = 1.0f-2, rtol = 1.0f-2)
+        if REACTANT_AVAILABLE
+            ps_ra, st_ra = (ps, st) |> xdev
+            u_ra, y_ra = (u, y) |> xdev
+
+            pred_ra = @jit deeponet((u_ra, y_ra), ps_ra, st_ra)
+            @test first(pred_ra) ≈ pred atol = 1.0f-2 rtol = 1.0f-2
+
+            @testset "check gradients" begin
+                (∂u_fd, ∂y_fd), ∂ps_fd = ∇sumabs2_finite_difference(
+                    deeponet, (u, y), ps, st
+                )
+                (∂u_ra, ∂y_ra), ∂ps_ra = ∇sumabs2_reactant(
+                    deeponet, (u_ra, y_ra), ps_ra, st_ra
+                )
+                (∂u_ra, ∂y_ra), ∂ps_ra = ((∂u_ra, ∂y_ra), ∂ps_ra) |> cpu_device()
+
+                @test ∂u_fd ≈ ∂u_ra atol = 1.0f-2 rtol = 1.0f-2
+                @test ∂y_fd ≈ ∂y_ra atol = 1.0f-2 rtol = 1.0f-2
+                @test check_approx(∂ps_fd, ∂ps_ra; atol = 1.0f-2, rtol = 1.0f-2)
+            end
         end
     end
 end

test/models/fno_tests.jl

@@ -3,6 +3,8 @@ using NeuralOperators, Test
 include("../shared_testsetup.jl")
 
 @testset "Fourier Neural Operator" begin
+    @test REACTANT_AVAILABLE || Sys.WORD_SIZE == 32
+
     rng = StableRNG(12345)
 
     @testset "complex finite-difference reference" begin
@@ -38,7 +40,9 @@ include("../shared_testsetup.jl")
         ),
     ]
 
-    xdev = reactant_device(; force = true)
+    if REACTANT_AVAILABLE
+        xdev = reactant_device(; force = true)
+    end
 
     @testset "$(length(setup.modes))D | shift=$(setup.shift)" for setup in setups
         fno = FourierNeuralOperator(; setup.chs, setup.modes, setup.shift)
@@ -50,22 +54,24 @@ include("../shared_testsetup.jl")
 
         @test size(first(fno(x, ps, st))) == setup.y_size
 
-        ps_ra, st_ra = (ps, st) |> xdev
-        x_ra, y_ra = (x, y) |> xdev
+        if REACTANT_AVAILABLE
+            ps_ra, st_ra = (ps, st) |> xdev
+            x_ra, y_ra = (x, y) |> xdev
 
-        res = first(fno(x, ps, st))
-        res_ra, _ = @jit fno(x_ra, ps_ra, st_ra)
-        @test res_ra ≈ res atol = 1.0f-2 rtol = 1.0f-2
+            res = first(fno(x, ps, st))
+            res_ra, _ = @jit fno(x_ra, ps_ra, st_ra)
+            @test res_ra ≈ res atol = 1.0f-2 rtol = 1.0f-2
 
-        @testset "check gradients" begin
-            ∂x_fd, ∂ps_fd = ∇sumabs2_finite_difference(fno, x, ps, st)
-            ∂x_ra, ∂ps_ra = ∇sumabs2_reactant(fno, x_ra, ps_ra, st_ra)
-            ∂x_ra, ∂ps_ra = (∂x_ra, ∂ps_ra) |> cpu_device()
+            @testset "check gradients" begin
+                ∂x_fd, ∂ps_fd = ∇sumabs2_finite_difference(fno, x, ps, st)
+                ∂x_ra, ∂ps_ra = ∇sumabs2_reactant(fno, x_ra, ps_ra, st_ra)
+                ∂x_ra, ∂ps_ra = (∂x_ra, ∂ps_ra) |> cpu_device()
 
-            # Float32 finite differencing is notoriously noisy, especially on Windows XLA.
-            # Relaxing tolerance to 1.0f-1 to prevent CI flakiness.
-            @test ∂x_fd ≈ ∂x_ra atol = 1.0f-1 rtol = 1.0f-1
-            @test check_approx(∂ps_fd, ∂ps_ra; atol = 1.0f-1, rtol = 1.0f-1)
+                # Float32 finite differencing is notoriously noisy, especially on Windows XLA.
+                # Relaxing tolerance to 1.0f-1 to prevent CI flakiness.
+                @test ∂x_fd ≈ ∂x_ra atol = 1.0f-1 rtol = 1.0f-1
+                @test check_approx(∂ps_fd, ∂ps_ra; atol = 1.0f-1, rtol = 1.0f-1)
+            end
         end
     end
 end

test/models/nomad_tests.jl

@@ -24,7 +24,9 @@ include("../shared_testsetup.jl")
         ),
     ]
 
-    xdev = reactant_device(; force = true)
+    if REACTANT_AVAILABLE
+        xdev = reactant_device(; force = true)
+    end
 
     @testset "$(setup.name)" for setup in setups
         u = rand(Float32, setup.u_size...)
@@ -36,20 +38,24 @@ include("../shared_testsetup.jl")
         pred = first(nomad((u, y), ps, st))
         @test setup.out_size == size(pred)
 
-        ps_ra, st_ra = xdev((ps, st))
-        u_ra, y_ra = xdev(u), xdev(y)
+        if REACTANT_AVAILABLE
+            ps_ra, st_ra = xdev((ps, st))
+            u_ra, y_ra = xdev(u), xdev(y)
 
-        pred_ra, _ = @jit nomad((u_ra, y_ra), ps_ra, st_ra)
-        @test pred_ra ≈ pred atol = 1.0f-2 rtol = 1.0f-2
+            pred_ra, _ = @jit nomad((u_ra, y_ra), ps_ra, st_ra)
+            @test pred_ra ≈ pred atol = 1.0f-2 rtol = 1.0f-2
 
-        @testset "check gradients" begin
-            (∂u_fd, ∂y_fd), ∂ps_fd = ∇sumabs2_finite_difference(nomad, (u, y), ps, st)
-            (∂u_ra, ∂y_ra), ∂ps_ra = ∇sumabs2_reactant(nomad, (u_ra, y_ra), ps_ra, st_ra)
-            (∂u_ra, ∂y_ra), ∂ps_ra = ((∂u_ra, ∂y_ra), ∂ps_ra) |> cpu_device()
+            @testset "check gradients" begin
+                (∂u_fd, ∂y_fd), ∂ps_fd = ∇sumabs2_finite_difference(nomad, (u, y), ps, st)
+                (∂u_ra, ∂y_ra), ∂ps_ra = ∇sumabs2_reactant(
+                    nomad, (u_ra, y_ra), ps_ra, st_ra
+                )
+                (∂u_ra, ∂y_ra), ∂ps_ra = ((∂u_ra, ∂y_ra), ∂ps_ra) |> cpu_device()
 
-            @test ∂u_fd ≈ ∂u_ra atol = 1.0f-2 rtol = 1.0f-2
-            @test ∂y_fd ≈ ∂y_ra atol = 1.0f-2 rtol = 1.0f-2
-            @test check_approx(∂ps_fd, ∂ps_ra; atol = 1.0f-2, rtol = 1.0f-2)
+                @test ∂u_fd ≈ ∂u_ra atol = 1.0f-2 rtol = 1.0f-2
+                @test ∂y_fd ≈ ∂y_ra atol = 1.0f-2 rtol = 1.0f-2
+                @test check_approx(∂ps_fd, ∂ps_ra; atol = 1.0f-2, rtol = 1.0f-2)
+            end
         end
     end
 end

test/shared_testsetup.jl

@@ -1,6 +1,10 @@
 using Lux, Optimisers, Random, StableRNGs, Reactant, Enzyme, FastTransforms
+using Reactant: Reactant_jll
 using MLDataDevices: cpu_device, reactant_device
 
+# Reactant_jll ships no 32-bit artifact.
+const REACTANT_AVAILABLE = Reactant_jll.is_available()
+
 sumabs2first(model, x, ps, st) = sum(abs2, first(model(x, ps, st)))
 
 # Local replacement for `LuxTestUtils.check_approx`: LuxTestUtils unconditionally skips