diff --git a/test/layers/layers_testsetup.jl b/test/layers/layers_testsetup.jl index 6c00268..0188648 100644 --- a/test/layers/layers_testsetup.jl +++ b/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 diff --git a/test/models/deeponet_tests.jl b/test/models/deeponet_tests.jl index 39e0f4b..56e3589 100644 --- a/test/models/deeponet_tests.jl +++ b/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 diff --git a/test/models/fno_tests.jl b/test/models/fno_tests.jl index 562ea3e..26f7418 100644 --- a/test/models/fno_tests.jl +++ b/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 diff --git a/test/models/nomad_tests.jl b/test/models/nomad_tests.jl index d24f362..063e6f5 100644 --- a/test/models/nomad_tests.jl +++ b/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 diff --git a/test/shared_testsetup.jl b/test/shared_testsetup.jl index a104629..43656d6 100644 --- a/test/shared_testsetup.jl +++ b/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