feat: update LuxTestUtils to support 1.12

avik-pal · avik-pal · commit 29d9a577dbbc · 2026-01-15T23:47:48.000-05:00
ci: test on 1.12

fix: ref negate

fix: use finite differences to test ground truth
diff --git a/.buildkite/testing_luxtestutils.yml b/.buildkite/testing_luxtestutils.yml
@@ -22,7 +22,7 @@ steps:
         matrix:
           setup:
             julia:
-              - "1.11"
+              - "1.12"
 
 env:
   SECRET_CODECOV_TOKEN: "jQ0BMTQgyZx7QGyU0Q2Ec7qB9mtE2q/tDu0FsfxvEG7/zOAGvXkyXrzIFFOQxvDoFcP+K2+hYZKMxicYdNqzr5wcxu505aNGN2GM3wyegAr+hO6q12bCFYx6qXzU9FLCCdeqINqn9gUSSOlGtWNFrbAlrTyz/D4Yo66TqBDzvaLL63FMnhCLaXW/zJt3hNuEAJaPY2O6Ze1rX2WZ3Y+i+s3uQ8aLImtoCJhPe8CRx+OhuYiTzGhynFfGntZ0738/1RN4gNM0S/hTC4gLE7XMVBanJpGh32rFaiDwW4zAyXKBrDkL3QA3MS1RvLTJxGJ085S16hCk0C4ddAhZCvIM9Q==;U2FsdGVkX1+bXdFeKMs5G79catOCyby2n07A2fg0FjVAvrjQLZ0yfvDS4paJiFikLkodho0khz2YALKb2Y0K6w=="
diff --git a/.github/workflows/CI_LuxTestUtils.yml b/.github/workflows/CI_LuxTestUtils.yml
@@ -19,5 +19,5 @@ jobs:
   test:
     uses: ./.github/workflows/CommonCI.yml
     with:
-      julia_version: "1.11"
+      julia_version: "1.12"
       project: "lib/LuxTestUtils"
diff --git a/Project.toml b/Project.toml
@@ -109,7 +109,7 @@ GPUArraysCore = "0.2"
 LinearAlgebra = "1.10"
 LossFunctions = "0.11.1, 1"
 LuxCore = "1.5.1"
-LuxLib = "1.15"
+LuxLib = "1.15.1"
 MLDataDevices = "1.17"
 MLUtils = "0.4.4"
 MPI = "0.20.19"
@@ -138,5 +138,5 @@ WeightInitializers = "1.3"
 Zygote = "0.7"
 julia = "1.10"
 
-[workspaces]
+[workspace]
 projects = ["test", "docs"]
diff --git a/lib/LuxCUDA/Project.toml b/lib/LuxCUDA/Project.toml
@@ -14,5 +14,5 @@ Reexport = "1"
 cuDNN = "1.4.3"
 julia = "1.10"
 
-[workspaces]
+[workspace]
 projects = ["test"]
diff --git a/lib/LuxCore/Project.toml b/lib/LuxCore/Project.toml
@@ -49,5 +49,5 @@ Setfield = "1"
 Tracker = "0.2.36"
 julia = "1.10"
 
-[workspaces]
+[workspace]
 projects = ["test"]
diff --git a/lib/LuxLib/Project.toml b/lib/LuxLib/Project.toml
@@ -105,5 +105,5 @@ UUIDs = "1.10"
 cuDNN = "1.3"
 julia = "1.10"
 
-[workspaces]
+[workspace]
 projects = ["test"]
diff --git a/lib/LuxLib/test/common_ops/dense_tests.jl b/lib/LuxLib/test/common_ops/dense_tests.jl
@@ -138,7 +138,7 @@ end
 end
 
 @testitem "Enzyme.Forward patch: dense" tags = [:common] setup = [SharedTestSetup] skip =
-    :(using LuxTestUtils; !LuxTestUtils.ENZYME_TESTING_ENABLED) begin
+    :(using LuxTestUtils; !LuxTestUtils.ENZYME_TESTING_ENABLED[]) begin
     using LuxLib, Random, ForwardDiff, Enzyme
 
     x = rand(Float32, 2, 2)
@@ -149,7 +149,7 @@ end
 end
 
 @testitem "Enzyme rules for fused dense" tags = [:common] setup = [SharedTestSetup] skip =
-    :(using LuxTestUtils; !LuxTestUtils.ENZYME_TESTING_ENABLED) begin
+    :(using LuxTestUtils; !LuxTestUtils.ENZYME_TESTING_ENABLED[]) begin
     using LuxLib, NNlib, Zygote, Enzyme
 
     # These are mostly for testing the CUDA rules since we don't enable the CUDA tests
diff --git a/lib/LuxLib/test/common_ops/dropout_tests.jl b/lib/LuxLib/test/common_ops/dropout_tests.jl
@@ -33,7 +33,8 @@
                 T(2),
                 dims;
                 atol=1.0f-3,
-                rtol=1.0f-3
+                rtol=1.0f-3,
+                ground_truth_eltype=Nothing
             )
 
             y, mask_, rng_ = dropout(rng, x, T(0.5), Val(false), T(2), dims)
@@ -86,7 +87,8 @@ end
                 T(2),
                 :;
                 atol=1.0f-3,
-                rtol=1.0f-3
+                rtol=1.0f-3,
+                ground_truth_eltype=Nothing
             )
 
             @jet sum(first(dropout(rng, x, mask, T(0.5), Val(true), Val(true), T(2), :)))
@@ -117,7 +119,8 @@ end
                 T(2),
                 :;
                 atol=1.0f-3,
-                rtol=1.0f-3
+                rtol=1.0f-3,
+                ground_truth_eltype=Nothing
             )
 
             @jet sum(first(dropout(rng, x, mask, T(0.5), Val(true), Val(false), T(2), :)))
@@ -170,7 +173,8 @@ end
                 T(0.5),
                 Val(true);
                 atol=1.0f-3,
-                rtol=1.0f-3
+                rtol=1.0f-3,
+                ground_truth_eltype=Nothing
             )
 
             @jet sum(first(alpha_dropout(rng, x, T(0.5), Val(true))))
diff --git a/lib/LuxTestUtils/Project.toml b/lib/LuxTestUtils/Project.toml
@@ -1,10 +1,11 @@
 name = "LuxTestUtils"
 uuid = "ac9de150-d08f-4546-94fb-7472b5760531"
-version = "2.1.0"
+version = "2.2.0"
 authors = ["Avik Pal <avikpal@mit.edu>"]
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
@@ -25,6 +26,7 @@ MLDataDevices = {path = "../MLDataDevices"}
 
 [compat]
 ADTypes = "1.10"
+Adapt = "4.1"
 ArrayInterface = "7.17.1"
 ChainRulesCore = "1.25.1"
 ComponentArrays = "0.15.22"
@@ -33,13 +35,13 @@ Enzyme = "0.13.81"
 FiniteDiff = "2.23.1"
 ForwardDiff = "0.10.36, 1"
 Functors = "0.5"
-JET = "0.9.6, 0.10"
+JET = "0.9.6, 0.10, 0.11"
 MLDataDevices = "1.17"
 Optimisers = "0.3.4, 0.4"
 Pkg = "1.10"
 Test = "1.10"
 Zygote = "0.7"
 julia = "1.10"
 
-[workspaces]
+[workspace]
 projects = ["test"]
diff --git a/lib/LuxTestUtils/src/LuxTestUtils.jl b/lib/LuxTestUtils/src/LuxTestUtils.jl
@@ -1,9 +1,10 @@
 module LuxTestUtils
 
+using Adapt: Adapt, adapt
 using ArrayInterface: ArrayInterface
 using ComponentArrays: ComponentArray, getdata, getaxes
 using DispatchDoctor: allow_unstable
-using Functors: Functors
+using Functors: Functors, fmap
 using MLDataDevices: cpu_device, gpu_device, get_device, get_device_type, AbstractGPUDevice
 using Optimisers: Optimisers
 using Pkg: PackageSpec
@@ -33,25 +34,41 @@ using Zygote: Zygote
 const CRC = ChainRulesCore
 const FD = FiniteDiff
 
+const JET_TESTING_ENABLED = Ref{Bool}(false)
+const ENZYME_TESTING_ENABLED = Ref{Bool}(false)
+const ZYGOTE_TESTING_ENABLED = Ref{Bool}(false)
+
 # Check if JET will work
 try
     using JET: JET, JETTestFailure, get_reports, report_call, report_opt
-    # XXX: In 1.11, JET leads to stack overflows
-    global JET_TESTING_ENABLED = v"1.10-" ≤ VERSION < v"1.11-"
+    JET_TESTING_ENABLED[] = true
 catch err
     @error "`JET.jl` did not successfully precompile on $(VERSION). All `@jet` tests will \
             be skipped." maxlog = 1 err = err
-    global JET_TESTING_ENABLED = false
+    JET_TESTING_ENABLED[] = false
 end
 
-# Check if Enzyme will work
-try
-    using Enzyme: Enzyme
-    __ftest(x) = x
-    Enzyme.autodiff(Enzyme.Reverse, __ftest, Enzyme.Active, Enzyme.Active(2.0))
-    global ENZYME_TESTING_ENABLED = Sys.islinux()
-catch err
-    global ENZYME_TESTING_ENABLED = false
+# Check if Enzyme will work (only on non-prerelease versions)
+@static if isempty(VERSION.prerelease)
+    try
+        using Enzyme: Enzyme
+        Enzyme.gradient(Enzyme.Reverse, Base.Fix1(sum, abs2), ones(Float32, 10))
+        ENZYME_TESTING_ENABLED[] = Sys.islinux()
+    catch err
+        @error "`Enzyme.jl` did not successfully differentiate a simple function or \
+                failed to load on $(VERSION). All Enzyme tests will be \
+                skipped." maxlog = 1 err = err
+        ENZYME_TESTING_ENABLED[] = false
+    end
+end
+
+function __init__()
+    ZYGOTE_TESTING_ENABLED[] = VERSION < v"1.12-"
+
+    if JET_TESTING_ENABLED[]
+        # JET doesn't work nicely on 1.11
+        JET_TESTING_ENABLED[] = VERSION < v"1.11-" || VERSION ≥ v"1.12-"
+    end
 end
 
 include("package_install.jl")
diff --git a/lib/LuxTestUtils/src/autodiff.jl b/lib/LuxTestUtils/src/autodiff.jl
@@ -1,18 +1,39 @@
+# taken from implementation in Lux.jl
+struct LuxEltypeAdaptor{T} end
+
+(l::LuxEltypeAdaptor)(x) = fmap(adapt(l), x)
+function (l::LuxEltypeAdaptor)(x::AbstractArray{T}) where {T}
+    return isbitstype(T) ? adapt(l, x) : map(adapt(l), x)
+end
+
+function Adapt.adapt_storage(
+    ::LuxEltypeAdaptor{T}, x::AbstractArray{<:AbstractFloat}
+) where {T<:AbstractFloat}
+    return convert(AbstractArray{T}, x)
+end
+
+function Adapt.adapt_storage(
+    ::LuxEltypeAdaptor{T}, x::AbstractArray{<:Complex{<:AbstractFloat}}
+) where {T<:AbstractFloat}
+    return convert(AbstractArray{Complex{T}}, x)
+end
+
 struct Constant{T}
     val::T
 end
 
-# Zygote.jl on CPU
-function ground_truth_gradient(f, args...)
+# FiniteDiff.jl on CPU
+function ground_truth_gradient(::Type{T}, backend, f, args...) where {T}
     cdev = cpu_device()
+    eltype_adaptor = T === Nothing ? identity : LuxEltypeAdaptor{T}()
     f_cpu = try
-        cdev(f)
+        eltype_adaptor(cdev(f))
     catch err
         @error "Encountered error while moving $(f) to CPU. Skipping movement... This can \
                 be fixed by defining overloads using ConstructionBase.jl" err
         f
     end
-    return gradient(f_cpu, AutoZygote(), map(cdev, args)...)
+    return gradient(f_cpu, backend, map(eltype_adaptor, map(cdev, args))...)
 end
 
 # Zygote.jl
@@ -31,8 +52,9 @@ function gradient(f::F, ::AutoEnzyme{Nothing}, args...) where {F}
 end
 
 function gradient(f::F, ad::AutoEnzyme{<:Enzyme.ReverseMode}, args...) where {F}
-    !ENZYME_TESTING_ENABLED &&
+    if !ENZYME_TESTING_ENABLED[]
         return ntuple(Returns(GradientComputationSkipped()), length(args))
+    end
 
     args_activity = map(args) do x
         needs_gradient(x) && return Enzyme.Duplicated(x, Enzyme.make_zero(x))
@@ -79,13 +101,13 @@ end
 """
     test_gradients(f, args...; skip_backends=[], broken_backends=[], kwargs...)
 
-Test the gradients of `f` with respect to `args` using the specified backends.
+Test the gradients of `f` with respect to `args` using the specified backends. The ground
+truth gradients are computed using FiniteDiff.jl (unless specified otherwise) on CPU.
 
 | Backend        | ADType              | CPU | GPU | Notes             |
 |:-------------- |:------------------- |:--- |:--- |:----------------- |
 | Zygote.jl      | `AutoZygote()`      | ✔   | ✔   |                   |
 | ForwardDiff.jl | `AutoForwardDiff()` | ✔   | ✖   | `len ≤ 32`        |
-| FiniteDiff.jl  | `AutoFiniteDiff()`  | ✔   | ✖   | `len ≤ 32`        |
 | Enzyme.jl      | `AutoEnzyme()`      | ✔   | ✖   | Only Reverse Mode |
 
 ## Arguments
@@ -105,6 +127,10 @@ Test the gradients of `f` with respect to `args` using the specified backends.
   - `enzyme_set_runtime_activity`: If `true`, then activate runtime activity for Enzyme.
   - `enable_enzyme_reverse_mode`: If `true`, then enable reverse mode for Enzyme.
   - `kwargs`: Additional keyword arguments to pass to `check_approx`.
+  - `ground_truth_backend`: The backend to use for computing the ground truth gradients.
+    Defaults to `AutoFiniteDiff()`.
+  - `ground_truth_eltype`: The eltype to use for computing the ground truth gradients.
+    Defaults to `Float64`.
 
 ## Example
 
@@ -125,6 +151,8 @@ function test_gradients(
     soft_fail::Union{Bool,Vector}=false,
     enzyme_set_runtime_activity::Bool=false,
     enable_enzyme_reverse_mode::Bool=false,
+    ground_truth_backend=AutoFiniteDiff(),
+    ground_truth_eltype::Union{Nothing,Type}=Float64,
     # Internal kwargs start
     source::LineNumberNode=LineNumberNode(0, nothing),
     test_expr::Expr=:(check_approx(∂args, ∂args_gt; kwargs...)),
@@ -156,9 +184,8 @@ function test_gradients(
     push!(backends, AutoZygote())
     if !on_gpu
         total_length ≤ 32 && push!(backends, AutoForwardDiff())
-        total_length ≤ 32 && push!(backends, AutoFiniteDiff())
         # TODO: Move Enzyme out of here once it supports GPUs
-        if enable_enzyme_reverse_mode || ENZYME_TESTING_ENABLED
+        if enable_enzyme_reverse_mode || ENZYME_TESTING_ENABLED[]
             mode = if enzyme_set_runtime_activity
                 Enzyme.set_runtime_activity(Enzyme.Reverse)
             else
@@ -177,18 +204,22 @@ function test_gradients(
     end
 
     # Test the gradients
-    ∂args_gt = ground_truth_gradient(f, args...)  # Should be Zygote in most cases
-
-    @assert (backends[1] ∉ broken_backends) && (backends[1] ∉ skip_backends) "first backend cannot be broken or skipped"
+    ∂args_gt = ground_truth_gradient(ground_truth_eltype, ground_truth_backend, f, args...)
+    backends = [
+        backend for backend in backends if
+        typeof(backend).name.name != typeof(ground_truth_backend).name.name
+    ]
 
     return @testset "gradtest($(f))" begin
         @testset "$(nameof(typeof(backend)))()" for backend in backends
             local_test_expr = :([$(nameof(typeof(backend)))] - $(test_expr))
 
             result = if check_ad_backend_in(backend, skip_backends)
                 Broken(:skipped, local_test_expr)
-            elseif (soft_fail isa Bool && soft_fail) ||
+            elseif (
+                (soft_fail isa Bool && soft_fail) ||
                 (soft_fail isa Vector && check_ad_backend_in(backend, soft_fail))
+            )
                 try
                     ∂args = allow_unstable() do
                         return gradient(f, backend, args...)
@@ -254,7 +285,7 @@ end
 
 removed_backend(::AutoTracker) = true
 removed_backend(::AutoReverseDiff) = true
-removed_backend(x) = false
+removed_backend(_) = false
 
 """
     @test_gradients(f, args...; kwargs...)
diff --git a/lib/LuxTestUtils/src/jet.jl b/lib/LuxTestUtils/src/jet.jl
@@ -53,7 +53,7 @@ Test Broken
 ```
 """
 macro jet(expr, args...)
-    !JET_TESTING_ENABLED && return :()
+    !JET_TESTING_ENABLED[] && return :()
 
     all_args, call_extras, opt_extras = [], [], []
     target_modules_set = false
diff --git a/lib/MLDataDevices/Project.toml b/lib/MLDataDevices/Project.toml
@@ -79,5 +79,5 @@ julia = "1.10"
 oneAPI = "1.5, 2"
 OpenCL = "0.10.5"
 
-[workspaces]
+[workspace]
 projects = ["test"]
diff --git a/lib/WeightInitializers/Project.toml b/lib/WeightInitializers/Project.toml
@@ -39,5 +39,5 @@ SpecialFunctions = "2.4"
 Statistics = "1.10"
 julia = "1.10"
 
-[workspaces]
+[workspace]
 projects = ["test"]
diff --git a/test/enzyme_tests.jl b/test/enzyme_tests.jl
@@ -127,7 +127,7 @@ end
             ps, st = dev(Lux.setup(rng, model))
             x = aType(x)
 
-            if LuxTestUtils.ENZYME_TESTING_ENABLED
+            if LuxTestUtils.ENZYME_TESTING_ENABLED[]
                 test_enzyme_gradients(model, x, ps, st)
             else
                 @test_broken false
@@ -154,7 +154,7 @@ end
             st = dev(st)
             x = aType(x)
 
-            if LuxTestUtils.ENZYME_TESTING_ENABLED
+            if LuxTestUtils.ENZYME_TESTING_ENABLED[]
                 test_enzyme_gradients(model, x, ps, st)
             else
                 @test_broken false
diff --git a/test/helpers/loss_tests.jl b/test/helpers/loss_tests.jl
diff --git a/test/helpers/training_tests.jl b/test/helpers/training_tests.jl
diff --git a/test/layers/basic_tests.jl b/test/layers/basic_tests.jl
diff --git a/test/layers/dropout_tests.jl b/test/layers/dropout_tests.jl

-Original file line number
+Diff line change
         matrix:
           setup:
             julia:
 -              - "1.11"
 +              - "1.12"
 env:
   SECRET_CODECOV_TOKEN: "jQ0BMTQgyZx7QGyU0Q2Ec7qB9mtE2q/tDu0FsfxvEG7/zOAGvXkyXrzIFFOQxvDoFcP+K2+hYZKMxicYdNqzr5wcxu505aNGN2GM3wyegAr+hO6q12bCFYx6qXzU9FLCCdeqINqn9gUSSOlGtWNFrbAlrTyz/D4Yo66TqBDzvaLL63FMnhCLaXW/zJt3hNuEAJaPY2O6Ze1rX2WZ3Y+i+s3uQ8aLImtoCJhPe8CRx+OhuYiTzGhynFfGntZ0738/1RN4gNM0S/hTC4gLE7XMVBanJpGh32rFaiDwW4zAyXKBrDkL3QA3MS1RvLTJxGJ085S16hCk0C4ddAhZCvIM9Q==;U2FsdGVkX1+bXdFeKMs5G79catOCyby2n07A2fg0FjVAvrjQLZ0yfvDS4paJiFikLkodho0khz2YALKb2Y0K6w=="