refactor: use setfield and make make_zero!! type-stable

ext/LuxEnzymeExt/LuxEnzymeExt.jl

@@ -3,6 +3,8 @@ module LuxEnzymeExt
 using ADTypes: AutoEnzyme
 using Enzyme: Enzyme, Active, Const, Duplicated
 using EnzymeCore: EnzymeCore
+using Setfield: @set!
+using Static: False, True
 
 using Lux: Lux
 using Lux.Training: TrainingBackendCache, TrainState

ext/LuxEnzymeExt/training.jl

@@ -1,42 +1,43 @@
 function Lux.Training.compute_gradients(
-        ::AutoEnzyme, obj_fn::F, data, ts::TrainState) where {F}
-    dps = Enzyme.make_zero(ts.parameters)
+        ad::AutoEnzyme, obj_fn::F, data, ts::TrainState) where {F}
+    dps = Lux.recursive_make_zero(ts.parameters)
 
     obj_fn_wrap, st_wrap, stats_wrap = Lux.Training.wrap_objective_function(
-        obj_fn, ts.model, ts.parameters, ts.states, data, Val(true))
+        obj_fn, ts.model, ts.parameters, ts.states, data, True())
 
     _, loss = Enzyme.autodiff(
         EnzymeCore.ReverseWithPrimal, Const(obj_fn_wrap), Active, Const(ts.model),
         Duplicated(ts.parameters, dps), Const(ts.states), Const(data))
 
-    cache = TrainingBackendCache{:Enzyme, false}(
-        dps, (; obj_fn=obj_fn_wrap, st_wrap, stats_wrap))
-    ts_new = TrainState(cache, obj_fn, ts.model, ts.parameters, st_wrap[],
-        ts.optimizer, ts.optimizer_state, ts.step)
-
-    return dps, loss, stats_wrap[], ts_new
+    cache = TrainingBackendCache(
+        ad, False(), dps, (; obj_fn=obj_fn_wrap, st_wrap, stats_wrap))
+    @set! ts.cache = cache
+    @set! ts.objective_function = obj_fn
+    @set! ts.states = st_wrap[]
+    return dps, loss, stats_wrap[], ts
 end
 
 const AUTODIFF_CACHE_TYPE = TrainingBackendCache{
-    :Enzyme, false, PS, <:NamedTuple{(:obj_fn, :st_wrap, :stats_wrap)}} where {PS}
+    <:AutoEnzyme, False, PS, <:NamedTuple{(:obj_fn, :st_wrap, :stats_wrap)}} where {PS}
 
 function Lux.Training.compute_gradients(
         ::AutoEnzyme, obj_fn::F, data, ts::TrainState{<:AUTODIFF_CACHE_TYPE, F}) where {F}
-    dps = Lux.recursive_make_zero!!(ts.cache.dparameters)
+    # dps = Lux.recursive_make_zero!!(ts.cache.dparameters)
+    Enzyme.make_zero!(ts.cache.dparameters)
+    dps = ts.cache.dparameters
 
     _, loss = Enzyme.autodiff(
         EnzymeCore.ReverseWithPrimal, Const(ts.cache.extras.obj_fn), Active,
         Const(ts.model), Duplicated(ts.parameters, dps), Const(ts.states), Const(data))
 
-    ts_new = TrainState(
-        ts.cache, obj_fn, ts.model, ts.parameters, ts.cache.extras.st_wrap[],
-        ts.optimizer, ts.optimizer_state, ts.step)
+    @set! ts.objective_function = obj_fn
+    @set! ts.states = ts.cache.extras.st_wrap[]
 
-    return dps, loss, ts.cache.extras.stats_wrap[], ts_new
+    return dps, loss, ts.cache.extras.stats_wrap[], ts
 end
 
 function Lux.Training.compute_gradients(ad::AutoEnzyme, obj_fn::F, data,
-        ts::TrainState{<:TrainingBackendCache{:Enzyme, false}}) where {F}
+        ts::TrainState{<:TrainingBackendCache{<:AutoEnzyme, False}}) where {F}
     @warn "Detected calls to `compute_gradients(::AutoEnzyme, ...)` with objective \
            function that is changing across function calls. This can lead to the \
            generation of slow code" maxlog=1
@@ -46,15 +47,14 @@ function Lux.Training.compute_gradients(ad::AutoEnzyme, obj_fn::F, data,
         Const{typeof(ts.model)}, Duplicated{typeof(ts.parameters)},
         Const{typeof(ts.states)}, Const{typeof(data)})
 
-    cache = TrainingBackendCache{:Enzyme, false}(ts.cache.dparameters, (; forward, reverse))
-    ts_new = TrainState(cache, obj_fn, ts.model, ts.parameters, ts.states,
-        ts.optimizer, ts.optimizer_state, ts.step)
-
-    return Lux.Training.compute_gradients(ad, obj_fn, data, ts_new)
+    cache = TrainingBackendCache(ad, False(), ts.cache.dparameters, (; forward, reverse))
+    @set! ts.cache = cache
+    @set! ts.objective_function = obj_fn
+    return Lux.Training.compute_gradients(ad, obj_fn, data, ts)
 end
 
 const AUTODIFF_THUNK_CACHE_TYPE = TrainingBackendCache{
-    :Enzyme, false, PS, <:NamedTuple{(:forward, :reverse)}} where {PS}
+    <:AutoEnzyme, False, PS, <:NamedTuple{(:forward, :reverse)}} where {PS}
 
 function Lux.Training.compute_gradients(::AutoEnzyme, obj_fn::F, data,
         ts::TrainState{<:AUTODIFF_THUNK_CACHE_TYPE, F}) where {F}
@@ -67,8 +67,7 @@ function Lux.Training.compute_gradients(::AutoEnzyme, obj_fn::F, data,
         Const(obj_fn), Const(ts.model), params, Const(ts.states), Const(data),
         (one(loss), Lux.recursive_make_zero(st_), Lux.recursive_make_zero(stats)), tape)
 
-    ts_new = TrainState(ts.cache, obj_fn, ts.model, ts.parameters, st_,
-        ts.optimizer, ts.optimizer_state, ts.step)
-
-    return dps, loss, stats, ts_new
+    @set! ts.objective_function = obj_fn
+    @set! ts.states = st_
+    return dps, loss, stats, ts
 end

ext/LuxReverseDiffExt/LuxReverseDiffExt.jl

@@ -5,9 +5,11 @@ using ArrayInterface: ArrayInterface
 using FunctionWrappers: FunctionWrapper
 using ReverseDiff: ReverseDiff, ForwardExecutor, ReverseExecutor, TrackedArray, TrackedReal,
                    @grad_from_chainrules
+using Setfield: @set!
+using Static: False, True
 
 using Lux: Lux, Utils
-using Lux.Training: TrainingBackendCache, TrainState
+using Lux.Training: Training, TrainingBackendCache, TrainState
 using LuxCore: LuxCore
 using MLDataDevices: CPUDevice
 

ext/LuxReverseDiffExt/training.jl

@@ -1,53 +1,50 @@
 # Uncompiled ReverseDiff
 function Lux.Training.compute_gradients(
         ad::AutoReverseDiff{false}, obj_fn::F, data, ts::TrainState) where {F}
-    grads = Lux.recursive_make_zero(ts.parameters)
-    ts_new = TrainState(
-        TrainingBackendCache{:ReverseDiff, true}(grads, nothing), obj_fn, ts.model,
-        ts.parameters, ts.states, ts.optimizer, ts.optimizer_state, ts.step)
-    return Lux.Training.compute_gradients(ad, obj_fn, data, ts_new)
+    @set! ts.cache = TrainingBackendCache(
+        ad, True(), Lux.recursive_make_zero(ts.parameters), nothing)
+    @set! ts.objective_function = obj_fn
+    return Lux.Training.compute_gradients(ad, obj_fn, data, ts)
 end
 
 function Lux.Training.compute_gradients(::AutoReverseDiff{false}, obj_fn::F, data,
-        ts::TrainState{<:TrainingBackendCache{:ReverseDiff, FT}}) where {F, FT}
-    dparams = FT ? ts.cache.dparameters : Lux.recursive_make_zero!!(ts.cache.dparameters)
+        ts::TrainState{<:TrainingBackendCache{AutoReverseDiff{false}}}) where {F}
+    dparams = Training.dparameters(ts.cache)
     tape = ReverseDiff.InstructionTape()
     ps_tracked = Lux.recursive_map(Utils.Fix3(TrackedArray, tape), ts.parameters, dparams)
 
     loss, st, stats = obj_fn(ts.model, ps_tracked, ts.states, data)
     loss.deriv = true
     ReverseDiff.reverse_pass!(tape)
 
-    ts_new = TrainState(
-        TrainingBackendCache{:ReverseDiff, false}(ts.cache.dparameters, nothing),
-        obj_fn, ts.model, ts.parameters, st, ts.optimizer, ts.optimizer_state, ts.step)
-
-    return ts.cache.dparameters, ReverseDiff.value(loss), stats, ts_new
+    @set! ts.cache.first_try = False()
+    @set! ts.objective_function = obj_fn
+    @set! ts.states = st
+    return dparams, ReverseDiff.value(loss), stats, ts
 end
 
 # Compiled ReverseDiff
 function Lux.Training.compute_gradients(
         ad::AutoReverseDiff{true}, obj_fn::F, data, ts::TrainState) where {F}
-    grads = Lux.recursive_make_zero(ts.parameters)
-    data_cache = deepcopy(data)
-    ps_cache = deepcopy(ts.parameters)
-    extras = (; data_cache, ps_cache)
-
-    ts_new = TrainState(
-        TrainingBackendCache{:ReverseDiff, true}(grads, extras), nothing, ts.model,
-        ts.parameters, ts.states, ts.optimizer, ts.optimizer_state, ts.step)
-    return Lux.Training.compute_gradients(ad, obj_fn, data, ts_new)
+    @set! ts.cache = TrainingBackendCache(
+        ad, True(), Lux.recursive_make_zero(ts.parameters),
+        (; data_cache=deepcopy(data), ps_cache=deepcopy(ts.parameters)))
+    @set! ts.objective_function = nothing
+
+    return Lux.Training.compute_gradients(ad, obj_fn, data, ts)
 end
 
 ## Tape hasn't been compiled yet / Function mismatch so recompile
-function Lux.Training.compute_gradients(::AutoReverseDiff{true}, obj_fn::F, data,
-        ts::TrainState{<:TrainingBackendCache{:ReverseDiff, FT}}) where {F, FT}
+function Lux.Training.compute_gradients(ad::AutoReverseDiff{true}, obj_fn::F, data,
+        ts::TrainState{<:TrainingBackendCache{AutoReverseDiff{true}}}) where {F}
     if LuxCore.statelength(ts.states) != 0
         throw(ArgumentError("AutoReverseDiff(; compile=true) is not supported for Lux \
                              models with non-empty state `st`."))
     end
 
-    if FT # do a dry run
+    first_try = ts.cache.first_try isa True
+
+    if first_try # do a dry run
         _, st_, stats = obj_fn(ts.model, ts.parameters, ts.states, data)
         if stats != NamedTuple()
             throw(ArgumentError("AutoReverseDiff(; compile=true) is not supported for \
@@ -59,20 +56,18 @@ function Lux.Training.compute_gradients(::AutoReverseDiff{true}, obj_fn::F, data
         end
     end
 
-    dparams = FT ? ts.cache.dparameters : Lux.recursive_make_zero!!(ts.cache.dparameters)
+    dparams = Training.dparameters(ts.cache)
 
     (; ps_cache, data_cache) = ts.cache.extras
-    if !FT
+    if !first_try
         Lux.recursive_copyto!(ps_cache, ts.parameters)
         Lux.recursive_copyto!(data_cache, data)
     end
 
-    obj_fn_wrap = first ∘ obj_fn
-
     tape = ReverseDiff.InstructionTape()
     ps_tracked = Lux.recursive_map(Utils.Fix3(TrackedArray, tape), ps_cache, dparams)
 
-    loss = obj_fn_wrap(ts.model, ps_tracked, ts.states, data_cache)
+    loss = first(obj_fn(ts.model, ps_tracked, ts.states, data_cache))
     loss.deriv = true
     ReverseDiff.reverse_pass!(tape)
 
@@ -81,18 +76,14 @@ function Lux.Training.compute_gradients(::AutoReverseDiff{true}, obj_fn::F, data
     reverse_executor = [FunctionWrapper{Nothing, Tuple{}}(ReverseExecutor(tape[i]))
                         for i in length(tape):-1:1]
 
-    compiled_extras = (;
-        ps_cache, data_cache, forward_executor, reverse_executor, output=loss)
-    ts_new = TrainState(
-        TrainingBackendCache{:ReverseDiff, false}(ts.cache.dparameters, compiled_extras),
-        obj_fn, ts.model, ts.parameters, ts.states,
-        ts.optimizer, ts.optimizer_state, ts.step)
-
-    return dparams, ReverseDiff.value(loss), NamedTuple(), ts_new
+    @set! ts.cache = TrainingBackendCache(ad, False(), dparams,
+        (; ps_cache, data_cache, forward_executor, reverse_executor, output=loss))
+    @set! ts.objective_function = obj_fn
+    return dparams, ReverseDiff.value(loss), NamedTuple(), ts
 end
 
 function Lux.Training.compute_gradients(::AutoReverseDiff{true}, obj_fn::F, data,
-        ts::TrainState{<:TrainingBackendCache{:ReverseDiff, false}, F}) where {F}
+        ts::TrainState{<:TrainingBackendCache{AutoReverseDiff{true}}, F}) where {F}
     (; ps_cache, data_cache, output) = ts.cache.extras
 
     dparams = Lux.recursive_make_zero!!(ts.cache.dparameters)
@@ -107,7 +98,7 @@ function Lux.Training.compute_gradients(::AutoReverseDiff{true}, obj_fn::F, data
         wrapper()
     end
 
-    ts_new = TrainState(ts.cache, obj_fn, ts.model, ts.parameters, ts.states,
-        ts.optimizer, ts.optimizer_state, ts.step)
-    return dparams, ReverseDiff.value(output), NamedTuple(), ts_new
+    @set! ts.cache.first_try = False()
+    @set! ts.objective_function = obj_fn
+    return dparams, ReverseDiff.value(output), NamedTuple(), ts
 end

ext/LuxTrackerExt/LuxTrackerExt.jl

@@ -3,10 +3,12 @@ module LuxTrackerExt
 using ADTypes: AbstractADType, AutoTracker
 using ArrayInterface: ArrayInterface
 using ChainRulesCore: ChainRulesCore
+using Setfield: @set!
+using Static: False, True
 using Tracker: Tracker, TrackedArray, TrackedReal, @grad_from_chainrules
 
 using Lux: Lux, Utils
-using Lux.Training: TrainingBackendCache, TrainState
+using Lux.Training: Training, TrainingBackendCache, TrainState
 
 const CRC = ChainRulesCore
 

ext/LuxTrackerExt/training.jl

@@ -1,23 +1,23 @@
 function Lux.Training.compute_gradients(::AutoTracker, obj_fn::F, data,
-        ts::TrainState{<:TrainingBackendCache{:Tracker, FT}}) where {F, FT}
-    dparams = FT ? ts.cache.dparameters : Lux.recursive_make_zero!!(ts.cache.dparameters)
-    ps_tracked = construct_tracked_params(ts.parameters, dparams)
+        ts::TrainState{<:TrainingBackendCache{AutoTracker}}) where {F}
+    dps = Training.dparameters(ts.cache)
+    ps_tracked = construct_tracked_params(ts.parameters, dps)
 
     loss, st, stats = obj_fn(ts.model, ps_tracked, ts.states, data)
     Tracker.back!(loss)
 
-    ts_new = TrainState(
-        TrainingBackendCache{:Tracker, false}(ts.cache.dparameters, nothing), obj_fn,
-        ts.model, ts.parameters, st, ts.optimizer, ts.optimizer_state, ts.step)
+    @set! ts.cache.first_try = False()
+    @set! ts.objective_function = obj_fn
+    @set! ts.states = st
 
-    return dparams, loss.data, stats, ts_new
+    return dps, loss.data, stats, ts
 end
 
 function Lux.Training.compute_gradients(
-        ::AutoTracker, obj_fn::F, data, ts::TrainState) where {F}
+        ad::AutoTracker, obj_fn::F, data, ts::TrainState) where {F}
     grads = Lux.recursive_make_zero(ts.parameters)
-    ts_new = TrainState(
-        TrainingBackendCache{:Tracker, true}(grads, nothing), obj_fn, ts.model,
-        ts.parameters, ts.states, ts.optimizer, ts.optimizer_state, ts.step)
-    return Lux.Training.compute_gradients(AutoTracker(), obj_fn, data, ts_new)
+    cache = TrainingBackendCache(ad, True(), grads, nothing)
+    @set! ts.cache = cache
+    @set! ts.objective_function = obj_fn
+    return Lux.Training.compute_gradients(ad, obj_fn, data, ts)
 end

src/helpers/recursive_ops.jl

@@ -102,18 +102,12 @@ function recursive_map(f::F, x::AbstractArray{T}, args...) where {F, T}
     (T <: Number || isbitstype(T)) && return f(x, args...) # Not all Number types (BigFloat) are bitstype
     return f.(x, args...)
 end
-function recursive_map(f::F, x::Tuple, args...) where {F}
+function recursive_map(f::F, x::Union{NamedTuple, Tuple}, args...) where {F}
     map_fn = let f = f
         (args_...) -> recursive_map(f, args_...)
     end
     return map(map_fn, x, args...)
 end
-function recursive_map(f::F, x::NamedTuple{fields}, args...) where {F, fields}
-    map_fn = let f = f
-        (args_...) -> recursive_map(f, args_...)
-    end
-    return NamedTuple{fields}(map(map_fn, values(x), values.(args)...))
-end
 recursive_map(f::F, x, args...) where {F} = fmap(f, x, args...)
 
 @compat(public,

src/helpers/training.jl

@@ -6,6 +6,7 @@ using ConcreteStructs: @concrete
 using FastClosures: @closure
 using Optimisers: Optimisers
 using Setfield: @set!
+using Static: StaticBool, Static, False, True
 
 using ..Lux: Lux
 using LuxCore: LuxCore, AbstractLuxLayer
@@ -50,13 +51,10 @@ Constructor for [`TrainState`](@ref).
 
 ## Arguments
 
-  - `rng`: Random Number Generator.
   - `ps`: Parameters of the model.
   - `st`: States of the model.
   - `model`: `Lux` model.
   - `optimizer`: Optimizer from `Optimisers.jl`.
-  - `transform_variables`: Function to transform the variables of the model. Typically used
-    to transfer variables to GPU / CPU.
 
 ## Returns
 
@@ -67,12 +65,18 @@ function TrainState(model::AbstractLuxLayer, ps, st, optimizer::Optimisers.Abstr
     return TrainState(nothing, nothing, model, ps, st, optimizer, st_opt, 0)
 end
 
-@concrete struct TrainingBackendCache{backend, first_try}
+@concrete struct TrainingBackendCache
+    backend
+    first_try <: StaticBool
     dparameters
     extras
 end
 
-training_backend(::TrainingBackendCache{backend}) where {backend} = backend
+dparameters(cache::TrainingBackendCache) = dparameters(cache, cache.first_try)
+function dparameters(cache::TrainingBackendCache, ::False)
+    return Lux.recursive_make_zero!!(cache.dparameters)
+end
+dparameters(cache::TrainingBackendCache, ::True) = cache.dparameters
 
 function Base.show(io::IO, ::MIME"text/plain", ts::TrainState)
     println(io, "TrainState")
@@ -83,8 +87,7 @@ function Base.show(io::IO, ::MIME"text/plain", ts::TrainState)
     print(io, "    step: ", ts.step)
     if ts.cache !== nothing
         if ts.cache isa TrainingBackendCache
-            print(io,
-                "\n    cache: $(nameof(typeof(ts.cache))){$(training_backend(ts.cache))}")
+            print(io, "\n    cache: $(nameof(typeof(ts.cache)))($(ts.cache.backend))")
         else
             print(io, "\n    cache: $(nameof(typeof(ts.cache)))")
         end
@@ -198,7 +201,7 @@ for package in (:Zygote, :Tracker, :ReverseDiff, :Enzyme)
     end
 end
 
-function generate_wrappers(::F, m, ps, st, data, ::Val{false}) where {F}
+function generate_wrappers(::F, m, ps, st, data, ::False) where {F}
     @warn "Detected function wrapper generation with function being updated between calls. \
            This will generate type-unstable code. A possible reason for this is \
            `TrainState` was compiled (first call to `compute_gradients`) with function \
@@ -208,13 +211,13 @@ function generate_wrappers(::F, m, ps, st, data, ::Val{false}) where {F}
 end
 
 # Run the code when trying to compile the function for the first time.
-function generate_wrappers(objective_function::F, m, ps, st, data, ::Val{true}) where {F}
+function generate_wrappers(objective_function::F, m, ps, st, data, ::True) where {F}
     _, stₙ, statsₙ = objective_function(m, ps, st, data)
     return Ref{typeof(stₙ)}(stₙ), Ref{typeof(statsₙ)}(statsₙ)
 end
 
 function wrap_objective_function(
-        objective_function::F, m, ps, st, data, first_try::Val) where {F}
+        objective_function::F, m, ps, st, data, first_try::StaticBool) where {F}
     st_updated, stats = generate_wrappers(objective_function, m, ps, st, data, first_try)
 
     wrapped_objective_function = @closure (model, ps, st, data) -> begin