refactor: use EnzymeRules.@easy_rule in Lux.jl (#1542)

* refactor: use `EnzymeRules.@easy_rule` in Lux.jl

* refactor: more migration

* fix: import Annotation

* fix: bump min Enzyme version

* fix: try some fixes

* feat: use enzyme rules for safe_reshape_internal

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.25.0"
+version = "1.26.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -87,7 +87,7 @@ ConcreteStructs = "0.2.3"
 DiffResults = "1.1"
 DispatchDoctor = "0.4.26"
 Enzyme = "0.13.81"
-EnzymeCore = "0.8.14"
+EnzymeCore = "0.8.15"
 FastClosures = "0.3.2"
 Flux = "0.16.3"
 ForwardDiff = "0.10.36, =1"

ext/LuxLossFunctionsExt.jl

@@ -44,46 +44,9 @@ function CRC.rrule(
     return LossFunctionImpl.fused_agg(sum, lfn, x, y), ∇fused_agg
 end
 
-# COV_EXCL_START
-
-function EnzymeRules.augmented_primal(
-    cfg::EnzymeRules.RevConfigWidth{1},
-    func::EnzymeCore.Const{typeof(LossFunctionImpl.fused_agg)},
-    ::Type{<:EnzymeCore.Active},
-    agg_f::EnzymeCore.Const{typeof(sum)},
-    lfn::EnzymeCore.Const{<:LossFunctions.Traits.Loss},
-    x::EnzymeCore.Annotation{<:AbstractArray},
-    y::EnzymeCore.Const,
-)
-    primal =
-        EnzymeRules.needs_primal(cfg) ? func.val(agg_f.val, lfn.val, x.val, y.val) : nothing
-
-    cache_x = EnzymeRules.overwritten(cfg)[4] ? copy(x.val) : nothing
-    cache_y = EnzymeRules.overwritten(cfg)[5] ? copy(y.val) : nothing
-
-    return EnzymeRules.AugmentedReturn(primal, nothing, (cache_x, cache_y))
-end
-
-function EnzymeRules.reverse(
-    cfg::EnzymeRules.RevConfigWidth{1},
-    ::EnzymeCore.Const{typeof(LossFunctionImpl.fused_agg)},
-    dret::EnzymeCore.Active,
-    (cache_x, cache_y),
-    agg_f::EnzymeCore.Const{typeof(sum)},
-    lfn::EnzymeCore.Const{<:LossFunctions.Traits.Loss},
-    x::EnzymeCore.Annotation{<:AbstractArray},
-    y::EnzymeCore.Const,
+EnzymeRules.@easy_rule(
+    LossFunctionImpl.fused_agg(fn::typeof(sum), lfn::LossFunctions.Traits.Loss, x, y),
+    (@Constant, @Constant, LossFunctions.deriv(lfn, x, y) .* Ω, @Constant)
 )
-    EnzymeRules.overwritten(cfg)[4] || (cache_x = x.val)
-    EnzymeRules.overwritten(cfg)[5] || (cache_y = y.val)
-
-    if !(typeof(x) <: EnzymeCore.Const)
-        @. x.dval = LossFunctions.deriv(lfn.val, cache_x, cache_y) * dret.val
-    end
-
-    return ntuple(Returns(nothing), 4)
-end
-
-# COV_EXCL_STOP
 
 end

lib/LuxLib/Project.toml

@@ -1,6 +1,6 @@
 name = "LuxLib"
 uuid = "82251201-b29d-42c6-8e01-566dec8acb11"
-version = "1.12.1"
+version = "1.13.0"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
 
 [deps]
@@ -74,7 +74,7 @@ CUDA = "5.8"
 ChainRulesCore = "1.25.1"
 DispatchDoctor = "0.4.12"
 Enzyme = "0.13.81"
-EnzymeCore = "0.8.14"
+EnzymeCore = "0.8.15"
 FastClosures = "0.3.2"
 ForwardDiff = "0.10.36, 1"
 Functors = "0.5"

lib/LuxLib/src/impl/Impl.jl

@@ -7,7 +7,7 @@ using StaticArraysCore: StaticVector, SArray
 using Static: StaticBool, True, False, static
 
 using ChainRulesCore: ChainRulesCore, NoTangent, HasReverseMode, RuleConfig
-using EnzymeCore: EnzymeCore, EnzymeRules
+using EnzymeCore: EnzymeCore, EnzymeRules, Annotation
 using ForwardDiff: ForwardDiff
 
 using KernelAbstractions: KernelAbstractions, @kernel, @Const, @index

lib/LuxLib/src/impl/batched_mul.jl

@@ -226,12 +226,7 @@ function batched_matmul_cpu!(
         batched_matmul_loopvec_impl!(z, x, y)
         return nothing
     end
-    if Utils.within_enzyme_autodiff()
-        # XXX: https://github.com/LuxDL/Lux.jl/issues/1024
-        fallback_batched_matmul!(z, LoopedArrayOp(), x, y)
-    else
-        NNlib.batched_mul!(z, x, y)
-    end
+    NNlib.batched_mul!(z, x, y)
     return nothing
 end
 

lib/LuxLib/src/impl/batchnorm.jl

@@ -87,10 +87,10 @@ function batchnorm_affine_normalize(
     β::Optional{<:AbstractVector},
     ϵ,
 ) where {F,xT,μT,σ²T,N}
-    x′ = reshape(x, :, size(x, N - 1), size(x, N))
-    return reshape(
+    x′ = Utils.safe_reshape(x, :, size(x, N - 1), size(x, N))
+    return Utils.safe_reshape(
         batchnorm_affine_normalize_internal(opmode, act, x′, vec(μ), vec(σ²), γ, β, ϵ),
-        size(x),
+        size(x)...,
     )
 end
 

lib/LuxLib/src/impl/groupnorm.jl

@@ -10,11 +10,15 @@ function groupnorm(
     act::F,
     ϵ,
 ) where {F,N,xT}
-    x′ = reshape(x, size(x)[1:(N - 2)]..., size(x, N - 1) ÷ groups, groups, size(x, N))
+    x′ = Utils.safe_reshape(
+        x, size(x)[1:(N - 2)]..., size(x, N - 1) ÷ groups, groups, size(x, N)
+    )
     (μ, σ²), _ = compute_batch_statistics(
         x′, nothing, nothing, groupnorm_reduce_dims(x), False(), nothing
     )
-    return reshape(groupnorm_affine_normalize(act, x′, μ, σ², γ, β, ϵ), size(x))
+    return Utils.safe_reshape(
+        groupnorm_affine_normalize(act, x′, μ, σ², γ, β, ϵ), size(x)...
+    )
 end
 
 function groupnorm_affine_normalize(
@@ -58,8 +62,8 @@ end
 ) where {F,N,xT,μT,σ²T}
     reshape_calls = if γ != Nothing
         quote
-            γ′ = reshape(γ, 1, size(x, N - 2), size(x, N - 1), 1)
-            β′ = reshape(β, 1, size(x, N - 2), size(x, N - 1), 1)
+            γ′ = Utils.safe_reshape(γ, 1, size(x, N - 2), size(x, N - 1), 1)
+            β′ = Utils.safe_reshape(β, 1, size(x, N - 2), size(x, N - 1), 1)
         end
     else
         quote
@@ -69,13 +73,13 @@ end
     end
 
     return quote
-        x′ = reshape(x, :, size(x, N - 2), size(x, N - 1), size(x, N))
-        μ′ = reshape(μ, 1, 1, size(x, N - 1), size(x, N))
-        σ²′ = reshape(σ², 1, 1, size(x, N - 1), size(x, N))
+        x′ = Utils.safe_reshape(x, :, size(x, N - 2), size(x, N - 1), size(x, N))
+        μ′ = Utils.safe_reshape(μ, 1, 1, size(x, N - 1), size(x, N))
+        σ²′ = Utils.safe_reshape(σ², 1, 1, size(x, N - 1), size(x, N))
         $(reshape_calls)
-        return reshape(
+        return Utils.safe_reshape(
             groupnorm_affine_normalize_internal(opmode, act, x′, μ′, σ²′, γ′, β′, ϵ),
-            size(x),
+            size(x)...,
         )
     end
 end

lib/LuxLib/src/impl/layernorm.jl

@@ -41,7 +41,7 @@ function expand_layernorm_dims(
 ) where {xT,γT,βT,N,M}
     new_γ_size = (size(γ)..., ntuple(i -> 1, N - M)...)
     new_β_size = (size(β)..., ntuple(i -> 1, N - M)...)
-    return reshape(γ, new_γ_size), reshape(β, new_β_size)
+    return Utils.safe_reshape(γ, new_γ_size...), Utils.safe_reshape(β, new_β_size...)
 end
 
 function expand_layernorm_dims(

lib/LuxLib/src/impl/normalization.jl

@@ -171,62 +171,8 @@ end
 
 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)
+    return Utils.safe_reshape(x, get_norm_reshape_dims(dims, length(x))...)
 end
-# COV_EXCL_STOP
 
 @inbounds function get_norm_reshape_dims(sx::NTuple{N,<:Int}, ly::Int) where {N}
     if ly == sx[N - 1]

lib/LuxLib/src/utils.jl

@@ -1,7 +1,7 @@
 module Utils
 
 using ChainRulesCore: ChainRulesCore
-using EnzymeCore: EnzymeCore, EnzymeRules
+using EnzymeCore: EnzymeCore, EnzymeRules, Annotation
 using FastClosures: @closure
 using Functors: Functors
 using ForwardDiff: ForwardDiff
@@ -40,9 +40,43 @@ ofeltype_array(::Type{T}, ::Nothing) where {T} = nothing
 contiguous(x::AbstractArray) = x
 contiguous(x::SubArray) = copy(x)
 
-safe_reshape(x::AbstractArray, dims...) = reshape(x, dims...)
+safe_reshape(x::AbstractArray, dims...) = safe_reshape_internal(x, dims)
 safe_reshape(::Nothing, dims...) = nothing
 
+safe_reshape_internal(x::AbstractArray, dims) = reshape(x, dims)
+
+# reshape sometimes causes trouble inside ReshapedArray for Enzyme. This is a workaround
+# for that.
+struct EnzymeRulesReshapeJacobian{T,N,D,A<:AbstractArray{T,N}} <: AbstractArray{T,2}
+    x::A
+    outsize::D
+end
+
+Base.size(J::EnzymeRulesReshapeJacobian) = (prod(J.outsize), length(J.x))
+
+function EnzymeRules.multiply_fwd_into(
+    previous, J::EnzymeRulesReshapeJacobian, x::AbstractArray
+)
+    reshaped_x = reshape(x, J.outsize...)
+    previous === nothing && return reshaped_x
+    @. previous += reshaped_x
+    return previous
+end
+
+function EnzymeRules.multiply_rev_into(
+    previous, J::EnzymeRulesReshapeJacobian, x::AbstractArray
+)
+    reshaped_x = reshape(x, size(J.x)...)
+    previous === nothing && return reshaped_x
+    @. previous += reshaped_x
+    return previous
+end
+
+EnzymeRules.@easy_rule(
+    safe_reshape_internal(x::AbstractArray, dims),
+    (EnzymeRulesReshapeJacobian(x, size(Ω)), @Constant),
+)
+
 remove_tracking(x) = x
 remove_tracking(x::AbstractArray) = x
 remove_tracking(::Type{T}) where {T} = T