feat: update LuxTestUtils to support 1.12

Author: avik-pal

lib/LuxTestUtils/Project.toml

@@ -1,7 +1,7 @@
 name = "LuxTestUtils"
 uuid = "ac9de150-d08f-4546-94fb-7472b5760531"
 authors = ["Avik Pal <[email protected]>"]
-version = "2.0.1"
+version = "2.1.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -32,7 +32,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.6.10"
 Optimisers = "0.3.4, 0.4"
 Test = "1.10"

lib/LuxTestUtils/src/LuxTestUtils.jl

@@ -32,25 +32,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("test_softfail.jl")

lib/LuxTestUtils/src/autodiff.jl

@@ -31,8 +31,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))
@@ -158,7 +159,7 @@ function test_gradients(
         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