test: enzyme tests are now passing (#1475)

avik-pal · web-flow · commit 1d0d9825b063 · 2025-09-12T18:18:20.000-04:00
* test: enzyme tests are now passing

* test: try switching off runtime activity

* chore: bump min enzyme version

* feat: add an enzyme rule to help out

* chore: remove dep

* fix: cleanup

* fix: dispatch

* fix: unwanted packages
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.21.2"
+version = "1.21.3"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -87,8 +87,8 @@ ComponentArrays = "0.15.28"
 ConcreteStructs = "0.2.3"
 DiffResults = "1.1"
 DispatchDoctor = "0.4.26"
-Enzyme = "0.13.49"
-EnzymeCore = "0.8.12"
+Enzyme = "0.13.74"
+EnzymeCore = "0.8.13"
 FastClosures = "0.3.2"
 Flux = "0.16.3"
 ForwardDiff = "0.10.36, 1"
diff --git a/docs/Project.toml b/docs/Project.toml
@@ -47,7 +47,7 @@ ComponentArrays = "0.15.22"
 Documenter = "1.4"
 DocumenterCitations = "1.3.6"
 DocumenterVitepress = "0.2"
-Enzyme = "0.13.35"
+Enzyme = "0.13.74"
 FiniteDiff = "2.23.1"
 Flux = "0.16.3"
 ForwardDiff = "0.10.36, 1"
diff --git a/examples/CIFAR10/Project.toml b/examples/CIFAR10/Project.toml
@@ -24,7 +24,7 @@ BFloat16s = "0.5.0"
 Comonicon = "1.0.8"
 ConcreteStructs = "0.2.3"
 DataAugmentation = "0.3"
-Enzyme = "0.13.35"
+Enzyme = "0.13.74"
 ImageCore = "0.10.2"
 ImageShow = "0.3.8"
 Interpolations = "0.15.1, 0.16"
diff --git a/examples/ConvolutionalVAE/Project.toml b/examples/ConvolutionalVAE/Project.toml
@@ -16,7 +16,7 @@ Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 [compat]
 ConcreteStructs = "0.2.3"
 DataAugmentation = "0.3.2"
-Enzyme = "0.13.35"
+Enzyme = "0.13.74"
 ImageShow = "0.3.8"
 Images = "0.26.1"
 Lux = "1.4.1"
diff --git a/examples/LSTMEncoderDecoder/Project.toml b/examples/LSTMEncoderDecoder/Project.toml
@@ -9,7 +9,7 @@ Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 
 [compat]
 CairoMakie = "0.13.6, 0.14, 0.15"
-Enzyme = "0.13.44"
+Enzyme = "0.13.74"
 Lux = "1.12.4"
 MLUtils = "0.4.8"
 Optimisers = "0.4.6"
diff --git a/examples/PINN2DPDE/Project.toml b/examples/PINN2DPDE/Project.toml
@@ -14,7 +14,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 [compat]
 ADTypes = "1.10"
 CairoMakie = "0.12.10, 0.13, 0.14, 0.15"
-Enzyme = "0.13"
+Enzyme = "0.13.74"
 Lux = "1"
 MLUtils = "0.4.4"
 OnlineStats = "1.7.1"
diff --git a/examples/RealNVP/Project.toml b/examples/RealNVP/Project.toml
@@ -13,7 +13,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 [compat]
 CairoMakie = "0.13.1, 0.14, 0.15"
 ConcreteStructs = "0.2.3"
-Enzyme = "0.13.35"
+Enzyme = "0.13.74"
 Lux = "1.5"
 MLUtils = "0.4.5"
 Optimisers = "0.4.6"
diff --git a/lib/LuxLib/Project.toml b/lib/LuxLib/Project.toml
@@ -1,7 +1,7 @@
 name = "LuxLib"
 uuid = "82251201-b29d-42c6-8e01-566dec8acb11"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.11.1"
+version = "1.11.2"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
@@ -72,8 +72,8 @@ CUDA = "5.8"
 ChainRulesCore = "1.24"
 Compat = "4.16"
 DispatchDoctor = "0.4.12"
-Enzyme = "0.13.35"
-EnzymeCore = "0.8.12"
+Enzyme = "0.13.74"
+EnzymeCore = "0.8.13"
 FastClosures = "0.3.2"
 ForwardDiff = "0.10.36, 1"
 Functors = "0.5"
diff --git a/lib/LuxLib/src/impl/batchnorm.jl b/lib/LuxLib/src/impl/batchnorm.jl
@@ -39,8 +39,8 @@ function batchnorm(
 ) where {F,xT,N}
     (μ, σ²), (rμ, rσ²) = compute_batch_statistics(
         x,
-        reshape_norm_dims(x, rμ),
-        reshape_norm_dims(x, rσ²),
+        reshape_norm_dims(rμ, size(x)),
+        reshape_norm_dims(rσ², size(x)),
         batchnorm_reduce_dims(x),
         training,
         momentum,
@@ -73,7 +73,7 @@ function batchnorm_affine_normalize(
     ϵ,
 ) where {F,xT,μT,σ²T,N}
     return affine_normalize(
-        act, x, μ, σ², reshape_norm_dims(x, γ), reshape_norm_dims(x, β), ϵ
+        act, x, μ, σ², reshape_norm_dims(γ, size(x)), reshape_norm_dims(β, size(x)), ϵ
     )
 end
 
diff --git a/lib/LuxLib/src/impl/groupnorm.jl b/lib/LuxLib/src/impl/groupnorm.jl
@@ -42,7 +42,7 @@ function groupnorm_affine_normalize(
     ϵ,
 ) where {F,N,xT,μT,σ²T}
     return affine_normalize(
-        act, x, μ, σ², reshape_norm_dims(x, γ), reshape_norm_dims(x, β), ϵ
+        act, x, μ, σ², reshape_norm_dims(γ, size(x)), reshape_norm_dims(β, size(x)), ϵ
     )
 end
 
diff --git a/lib/LuxLib/src/impl/normalization.jl b/lib/LuxLib/src/impl/normalization.jl
@@ -159,18 +159,74 @@ function normalization(
 ) where {F}
     (μ, σ²), (rμ, rσ²) = compute_batch_statistics(
         x,
-        reshape_norm_dims(x, rμ),
-        reshape_norm_dims(x, rσ²),
+        reshape_norm_dims(rμ, size(x)),
+        reshape_norm_dims(rσ², size(x)),
         reduce_dims,
         training,
         momentum,
     )
-    γ, β = reshape_norm_dims(x, γ), reshape_norm_dims(x, β)
+    γ, β = reshape_norm_dims(γ, size(x)), reshape_norm_dims(β, size(x))
     return affine_normalize(act, x, μ, σ², γ, β, epsilon), rμ, rσ²
 end
 
-reshape_norm_dims(_, ::Nothing) = nothing
-reshape_norm_dims(y, x) = reshape(x, get_norm_reshape_dims(size(y), length(x)))
+reshape_norm_dims(::Nothing, ::Dims) = nothing
+function reshape_norm_dims(x::AbstractArray, dims::Dims)
+    if Utils.within_enzyme_autodiff()
+        y = similar(x, get_norm_reshape_dims(dims, length(x)))
+        reshape_norm_dims!(y, x)
+        return y
+    else
+        return reshape(x, get_norm_reshape_dims(dims, length(x)))
+    end
+end
+
+function reshape_norm_dims!(y::AbstractArray, x::AbstractArray)
+    copyto!(vec(y), vec(x))
+    return nothing
+end
+
+function CRC.rrule(::typeof(reshape_norm_dims), x::AbstractArray, dims::Dims)
+    y = reshape_norm_dims(x, dims)
+    ∇reshape_norm_dims = @closure Δ -> begin
+        ∂x = CRC.@thunk reshape(recursive_unthunk(Δ), size(x))
+        return ∂∅, ∂x, ∂∅
+    end
+    return y, ∇reshape_norm_dims
+end
+
+# COV_EXCL_START
+# reshape_norm_dims is a constant source of runtime activity for Enzyme. Define custom
+# rules to avoid this.
+function EnzymeRules.augmented_primal(
+    cfg::EnzymeRules.RevConfigWidth{1},
+    ::EnzymeCore.Const{typeof(reshape_norm_dims)},
+    ::Type{EnzymeCore.Const{Nothing}},
+    y::EnzymeCore.Annotation{<:AbstractArray},
+    x::EnzymeCore.Annotation{<:AbstractArray},
+)
+    if EnzymeRules.needs_primal(cfg)
+        copyto!(vec(y.val), vec(x.val))
+    end
+    return EnzymeRules.AugmentedReturn(nothing, nothing, ())
+end
+
+function EnzymeRules.reverse(
+    ::EnzymeRules.RevConfigWidth{1},
+    ::EnzymeCore.Const{typeof(reshape_norm_dims)},
+    ::Type{EnzymeCore.Const{Nothing}},
+    tape,
+    y::EnzymeCore.Annotation{<:AbstractArray},
+    x::EnzymeCore.Annotation{<:AbstractArray},
+)
+    if !(typeof(y) <: EnzymeCore.Const)
+        if !(typeof(x) <: EnzymeCore.Const)
+            copyto!(vec(x.dval), vec(y.dval))
+        end
+        fill!(y.dval, false)
+    end
+    return ntuple(Returns(nothing), 2)
+end
+# COV_EXCL_STOP
 
 @inbounds function get_norm_reshape_dims(sx::NTuple{N,<:Int}, ly::Int) where {N}
     if ly == sx[N - 1]
diff --git a/lib/LuxLib/test/Project.toml b/lib/LuxLib/test/Project.toml
@@ -41,7 +41,7 @@ BenchmarkTools = "1.5"
 CPUSummary = "0.2.6"
 ChainRulesCore = "1.24"
 ComponentArrays = "0.15.22"
-Enzyme = "0.13.35"
+Enzyme = "0.13.74"
 EnzymeCore = "0.8.12"
 ExplicitImports = "1.9.0"
 ForwardDiff = "0.10.36, 1"
diff --git a/lib/LuxLib/test/common_ops/dropout_tests.jl b/lib/LuxLib/test/common_ops/dropout_tests.jl
@@ -116,7 +116,6 @@ end
                 Val(false),
                 T(2),
                 :;
-                broken_backends=length(x_shape) > 2 ? [AutoEnzyme()] : [],
                 atol=1.0f-3,
                 rtol=1.0f-3
             )
diff --git a/lib/LuxLib/test/normalization/batchnorm_tests.jl b/lib/LuxLib/test/normalization/batchnorm_tests.jl
@@ -103,7 +103,6 @@ function run_batchnorm_testing(gen_f, T, sz, training, affine, track_stats, act,
     end
 
     if is_training(training)
-        # XXX: Fails due to runtime activity but setting it doesn't help
         @test_gradients(
             sumabs2first,
             batchnorm,
@@ -119,8 +118,6 @@ function run_batchnorm_testing(gen_f, T, sz, training, affine, track_stats, act,
             atol,
             rtol,
             soft_fail=[AutoFiniteDiff()],
-            skip_backends=[AutoEnzyme()],
-            enzyme_set_runtime_activity=true
         )
     end
 end
@@ -202,8 +199,7 @@ end
             0.9,
             1.0e-5;
             atol=1.0f-3,
-            rtol=1.0f-3,
-            broken_backends=[AutoEnzyme()]
+            rtol=1.0f-3
         )
     end
 end
diff --git a/lib/LuxLib/test/normalization/instancenorm_tests.jl b/lib/LuxLib/test/normalization/instancenorm_tests.jl
@@ -42,7 +42,6 @@ function run_instancenorm_testing(gen_f, T, sz, training, act, aType)
             atol,
             rtol,
             soft_fail=[AutoFiniteDiff()],
-            enzyme_set_runtime_activity=true
         )
     end
 
@@ -75,7 +74,6 @@ function run_instancenorm_testing(gen_f, T, sz, training, act, aType)
             atol,
             rtol,
             soft_fail=[AutoFiniteDiff()],
-            enzyme_set_runtime_activity=true
         )
     end
 end
diff --git a/lib/LuxTestUtils/Project.toml b/lib/LuxTestUtils/Project.toml
@@ -28,7 +28,7 @@ ArrayInterface = "7.17.1"
 ChainRulesCore = "1.24.0"
 ComponentArrays = "0.15.22"
 DispatchDoctor = "0.4.12"
-Enzyme = "0.13.35"
+Enzyme = "0.13.74"
 FiniteDiff = "2.23.1"
 ForwardDiff = "0.10.36, 1"
 Functors = "0.5"
diff --git a/perf/resnet/Project.toml b/perf/resnet/Project.toml
@@ -16,7 +16,7 @@ LuxCUDA = {path = "../../lib/LuxCUDA"}
 [compat]
 BenchmarkTools = "1.6.0"
 Comonicon = "1.0.8"
-Enzyme = "0.13.46"
+Enzyme = "0.13.74"
 Lux = "1.13.3"
 LuxCUDA = "0.3.3"
 Random = "1.11"
diff --git a/test/Project.toml b/test/Project.toml
@@ -56,7 +56,7 @@ ChainRulesCore = "1.24"
 ComponentArrays = "0.15.22"
 DispatchDoctor = "0.4.12"
 Documenter = "1.4"
-Enzyme = "0.13.35"
+Enzyme = "0.13.74"
 ExplicitImports = "1.9.0"
 ForwardDiff = "0.10.36, 1"
 Functors = "0.5"
diff --git a/test/contrib/freeze_tests.jl b/test/contrib/freeze_tests.jl
@@ -29,16 +29,7 @@
 
             @test m(x, ps, st)[1] == m(x, ps_c, st)[1]
             @jet m(x, ps_c, st)
-            @test_gradients(
-                sumabs2first,
-                m,
-                x,
-                ps_c,
-                st;
-                atol=1.0f-3,
-                rtol=1.0f-3,
-                enzyme_set_runtime_activity=true
-            )
+            @test_gradients(sumabs2first, m, x, ps_c, st; atol=1.0f-3, rtol=1.0f-3,)
         end
 
         @testset "LuxDL/Lux.jl#427" begin
@@ -85,16 +76,7 @@ end
 
         @test d(x, psd, std)[1] == fd(x, ps, st)[1]
         @jet fd(x, ps, st)
-        @test_gradients(
-            sumabs2first,
-            fd,
-            x,
-            ps,
-            st;
-            atol=1.0f-3,
-            rtol=1.0f-3,
-            enzyme_set_runtime_activity=true
-        )
+        @test_gradients(sumabs2first, fd, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
 
         fd = Lux.Experimental.freeze(d, ())
         @test fd === d
diff --git a/test/helpers/loss_tests.jl b/test/helpers/loss_tests.jl
@@ -216,13 +216,7 @@ end
             @jet bceloss(σ.(logŷ), y)
             @jet bceloss_smooth(σ.(logŷ), y)
 
-            @test_gradients(
-                Base.Fix2(bceloss, y),
-                σ.(logŷ);
-                atol=1.0f-3,
-                rtol=1.0f-3,
-                enzyme_set_runtime_activity=true
-            )
+            @test_gradients(Base.Fix2(bceloss, y), σ.(logŷ); atol=1.0f-3, rtol=1.0f-3,)
         end
 
         @testset "Logit BinaryCrossEntropyLoss" begin
@@ -242,13 +236,7 @@ end
             @jet logitbceloss(logŷ, y)
             @jet logitbceloss_smooth(logŷ, y)
 
-            @test_gradients(
-                Base.Fix2(logitbceloss, y),
-                logŷ;
-                atol=1.0f-3,
-                rtol=1.0f-3,
-                enzyme_set_runtime_activity=true
-            )
+            @test_gradients(Base.Fix2(logitbceloss, y), logŷ; atol=1.0f-3, rtol=1.0f-3,)
         end
 
         @testset "BinaryFocalLoss" begin
diff --git a/test/helpers/training_tests.jl b/test/helpers/training_tests.jl
@@ -82,8 +82,6 @@ end
                 continue
             !LuxTestUtils.ENZYME_TESTING_ENABLED && ad isa AutoEnzyme && continue
 
-            broken = ad isa AutoEnzyme && VERSION ≥ v"1.11-"
-
             ps, st = dev(Lux.setup(rng, model))
             tstate = Training.TrainState(model, ps, st, opt)
 
@@ -116,7 +114,7 @@ end
                 )
 
                 final_loss * 100 < initial_loss
-            end broken = broken
+            end
 
             # Test the adjust API
             tstate = Optimisers.adjust(tstate, 0.1f0)
diff --git a/test/layers/attention_tests.jl b/test/layers/attention_tests.jl
diff --git a/test/layers/normalize_tests.jl b/test/layers/normalize_tests.jl
diff --git a/test/runtests.jl b/test/runtests.jl

Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,7 @@ function groupnorm_affine_normalize(`
`42`	`42`	`ϵ,`
`43`	`43`	`) where {F,N,xT,μT,σ²T}`
`44`	`44`	`return affine_normalize(`
`45`		`- act, x, μ, σ², reshape_norm_dims(x, γ), reshape_norm_dims(x, β), ϵ`
	`45`	`+ act, x, μ, σ², reshape_norm_dims(γ, size(x)), reshape_norm_dims(β, size(x)), ϵ`
`46`	`46`	`)`
`47`	`47`	`end`
`48`	`48`