feat: update LuxTestUtils to support 1.12

ci: test on 1.12

fix: ref negate

fix: use finite differences to test ground truth

Author: avik-pal

.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=="

.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"

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

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,7 +35,7 @@ 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"

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")

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
+# FiniteDiff.jl on CPU
 function ground_truth_gradient(f, args...)
     cdev = cpu_device()
+    f64 = LuxEltypeAdaptor{Float64}()
     f_cpu = try
-        cdev(f)
+        f64(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, AutoFiniteDiff(), map(f64, 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 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
@@ -156,9 +178,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

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

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

test/helpers/loss_tests.jl

@@ -22,7 +22,7 @@
     ∂x1 = ForwardDiff.derivative(Base.Fix2(LuxOps.xlogy, 3.0), 2.0)
     ∂y1 = ForwardDiff.derivative(Base.Fix1(LuxOps.xlogy, 2.0), 3.0)
     ∂x2, ∂y2 = Zygote.gradient(LuxOps.xlogy, 2.0, 3.0)
-    if LuxTestUtils.ENZYME_TESTING_ENABLED
+    if LuxTestUtils.ENZYME_TESTING_ENABLED[]
         ((∂x3, ∂y3),) = Enzyme.autodiff(
             Enzyme.Reverse, LuxOps.xlogy, Active, Active(2.0), Active(3.0)
         )
@@ -38,7 +38,7 @@
     @test @inferred(LuxOps.xlogy(0, 1)) isa Number
     @jet LuxOps.xlogy(2, 3)
 
-    if LuxTestUtils.ENZYME_TESTING_ENABLED
+    if LuxTestUtils.ENZYME_TESTING_ENABLED[]
         @test @inferred(
             Enzyme.autodiff(Enzyme.Reverse, LuxOps.xlogy, Active, Active(2.0), Active(3.0))
         ) isa Any

test/helpers/training_tests.jl

@@ -40,7 +40,7 @@ end
 
         for ad in (AutoZygote(), AutoTracker(), AutoReverseDiff(), AutoEnzyme())
             ongpu && (ad isa AutoReverseDiff || ad isa AutoEnzyme) && continue
-            !LuxTestUtils.ENZYME_TESTING_ENABLED && ad isa AutoEnzyme && continue
+            !LuxTestUtils.ENZYME_TESTING_ENABLED[] && ad isa AutoEnzyme && continue
 
             grads, _, _, _ = Training.compute_gradients(ad, _loss_function, x, tstate)
             tstate_ = Training.apply_gradients(tstate, grads)
@@ -80,7 +80,7 @@ end
             ongpu &&
                 (ad isa AutoReverseDiff || ad isa AutoEnzyme || ad isa AutoMooncake) &&
                 continue
-            !LuxTestUtils.ENZYME_TESTING_ENABLED && ad isa AutoEnzyme && continue
+            !LuxTestUtils.ENZYME_TESTING_ENABLED[] && ad isa AutoEnzyme && continue
 
             function get_total_loss(model, tstate)
                 loss = 0.0f0
@@ -215,7 +215,7 @@ end
 end
 
 @testitem "Training API Enzyme Runtime Mode" setup = [SharedTestSetup] tags = [:misc] skip =
-    :(using LuxTestUtils; !LuxTestUtils.ENZYME_TESTING_ENABLED) begin
+    :(using LuxTestUtils; !LuxTestUtils.ENZYME_TESTING_ENABLED[]) begin
     using Lux, Random, Enzyme, Optimisers
 
     function makemodel(n)
@@ -264,7 +264,7 @@ end
 
 @testitem "Enzyme: Invalidate Cache on State Update" setup = [SharedTestSetup] tags = [
     :misc
-] skip = :(using LuxTestUtils; !LuxTestUtils.ENZYME_TESTING_ENABLED) begin
+] skip = :(using LuxTestUtils; !LuxTestUtils.ENZYME_TESTING_ENABLED[]) begin
     using ADTypes, Optimisers
 
     mse = MSELoss()