feat: update minimum version of Enzyme to 0.13

LuxDL · Sep 21, 2024 · 6dd7701 · 6dd7701
1 parent 79ed8fe
commit 6dd7701
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 10 deletions.
diff --git a/Project.toml b/Project.toml
@@ -59,7 +59,7 @@ ChainRulesCore = "1.24"
 Compat = "4.15.0"
 CpuId = "0.3"
 DispatchDoctor = "0.4.12"
-EnzymeCore = "0.7.7"
+EnzymeCore = "0.8"
 FastClosures = "0.3.2"
 ForwardDiff = "0.10.36"
 Hwloc = "3.2"

diff --git a/src/impl/activation.jl b/src/impl/activation.jl
@@ -213,18 +213,19 @@ end
 # Enzyme works for all of these except `gelu`.
 # See https://github.com/EnzymeAD/Enzyme.jl/issues/1671
 function EnzymeRules.augmented_primal(
-        cfg::EnzymeRules.ConfigWidth{1}, func::EnzymeCore.Const{typeof(gelu)},
+        cfg::EnzymeRules.RevConfigWidth{1}, func::EnzymeCore.Const{typeof(gelu)},
         ::Type{<:EnzymeCore.Active}, x::EnzymeCore.Active{<:Number})
     primal = EnzymeRules.needs_primal(cfg) ? func.val(x.val) : nothing
     return EnzymeRules.AugmentedReturn(primal, nothing, nothing)
 end
 
 function EnzymeRules.reverse(
-        ::EnzymeRules.ConfigWidth{1}, ::EnzymeCore.Const{typeof(gelu)},
+        ::EnzymeRules.RevConfigWidth{1}, ::EnzymeCore.Const{typeof(gelu)},
         dret::EnzymeCore.Active, ::Nothing, x::EnzymeCore.Active{<:Number})
     return (dret.val * ∇gelu(x.val),)
 end
 
+# FIXME: ForwardRules changed in EnzymeCore 0.8
 function EnzymeRules.forward(
         ::EnzymeCore.Const{typeof(gelu)}, ::Type{<:EnzymeCore.Duplicated},
         x::EnzymeCore.Duplicated{<:Number})

diff --git a/src/impl/batched_mul.jl b/src/impl/batched_mul.jl
@@ -133,7 +133,7 @@ end
 for func in (NNlib.batched_mul!, batched_matmul_loopvec_impl!)
     @eval begin
         function EnzymeRules.augmented_primal(
-                cfg::EnzymeRules.ConfigWidth, ::EnzymeCore.Const{typeof($(func))},
+                cfg::EnzymeRules.RevConfigWidth, ::EnzymeCore.Const{typeof($(func))},
                 ::Type{RT}, C::EnzymeCore.Annotation{<:AbstractArray{<:Any, 3}},
                 A::EnzymeCore.Annotation{<:AbstractArray{<:Any, 3}},
                 B::EnzymeCore.Annotation{<:AbstractArray{<:Any, 3}}) where {RT}
@@ -155,7 +155,7 @@ for func in (NNlib.batched_mul!, batched_matmul_loopvec_impl!)
         end
 
         function EnzymeRules.reverse(
-                cfg::EnzymeRules.ConfigWidth, ::EnzymeCore.Const{typeof($(func))},
+                cfg::EnzymeRules.RevConfigWidth, ::EnzymeCore.Const{typeof($(func))},
                 ::Type{RT}, cache, C::EnzymeCore.Annotation{<:AbstractArray{<:Any, 3}},
                 A::EnzymeCore.Annotation{<:AbstractArray{<:Any, 3}},
                 B::EnzymeCore.Annotation{<:AbstractArray{<:Any, 3}}) where {RT}

diff --git a/src/utils.jl b/src/utils.jl
@@ -233,7 +233,7 @@ CRC.@non_differentiable safe_minimum(::Any...)
 macro enzyme_alternative(f₁, f₂)
     return esc(quote
         function EnzymeRules.augmented_primal(
-                ::EnzymeRules.ConfigWidth, ::EnzymeCore.Const{typeof($(f₁))},
+                ::EnzymeRules.RevConfig, ::EnzymeCore.Const{typeof($(f₁))},
                 ::Type{RT}, args...) where {RT}
             fwd, rev = EnzymeCore.autodiff_thunk(
                 EnzymeCore.ReverseSplitWithPrimal, EnzymeCore.Const{typeof($(f₂))},
@@ -245,11 +245,12 @@ macro enzyme_alternative(f₁, f₂)
         end
 
         function EnzymeRules.reverse(
-                ::EnzymeRules.ConfigWidth, ::EnzymeCore.Const{typeof($(f₁))},
+                ::EnzymeRules.RevConfig, ::EnzymeCore.Const{typeof($(f₁))},
                 ::Type{RT}, (tape, rev), args...) where {RT}
             return only(rev(EnzymeCore.Const($(f₂)), args..., tape))
         end
 
+        # FIXME: ForwardRules changed in EnzymeCore 0.8
         function EnzymeRules.forward(
                 ::EnzymeCore.Const{typeof($(f₁))}, ::Type{RT}, args...) where {RT}
             EnzymeCore.autodiff(EnzymeCore.Forward, EnzymeCore.Const($(f₂)), RT, args...)

diff --git a/test/Project.toml b/test/Project.toml
@@ -37,14 +37,14 @@ BLISBLAS = "0.1"
 BenchmarkTools = "1.5"
 ChainRulesCore = "1.24"
 ComponentArrays = "0.15.16"
-Enzyme = "0.12.26"
-EnzymeCore = "0.7.7"
+Enzyme = "0.13.1"
+EnzymeCore = "0.8"
 ExplicitImports = "1.9.0"
 ForwardDiff = "0.10.36"
 Hwloc = "3.2"
 InteractiveUtils = "<0.0.1, 1"
 JLArrays = "0.1.5"
-LuxTestUtils = "1.2"
+LuxTestUtils = "1.2.1"
 MKL = "0.7"
 MLDataDevices = "1.0.0"
 NNlib = "0.9.21"