fix: offload all pathces to Optimisers

Author: avik-pal

ext/LuxReactantExt/LuxReactantExt.jl

@@ -2,7 +2,8 @@ module LuxReactantExt
 
 using Enzyme: Enzyme, Const
 using Optimisers: Optimisers
-using Reactant: Reactant, @compile, @code_hlo, AnyTracedRArray, TracedRArray, TracedRNumber
+using Reactant:
+    Reactant, @compile, @code_hlo, @jit, AnyTracedRArray, TracedRArray, TracedRNumber
 using ReactantCore: ReactantCore, @trace
 using Setfield: @set!
 using Static: True, False

ext/LuxReactantExt/patches.jl

@@ -3,20 +3,7 @@ Utils.vec(x::AnyTracedRArray) = Reactant.TracedUtils.materialize_traced_array(ve
 # XXX: Use PoolDims once EnzymeJAX supports stablehlo.reduce_window adjoint
 Lux.calculate_pool_dims(g::Lux.GlobalPoolMode, ::TracedRArray) = g
 
-# Optimisers AccumGrad
-function Optimisers.apply!(opt::Lux.ReactantCompatibleOptimisers.AccumGrad, state, x, dx)
-    accum_dx, counter = state
-    @trace if counter == 1
-        @. accum_dx = dx / opt.n
-    else
-        @. accum_dx += dx / opt.n
-    end
-    @trace if counter == opt.n
-        dx_final = dx
-        counter = 1
-    else
-        dx_final = zero.(dx)
-        counter += 1
-    end
-    return (accum_dx, counter), dx_final
+# Optimisers setup
+function Lux.Training.optimisers_setup_with_jit(opt, ps)
+    return @jit Optimisers.setup(opt, ps)
 end

src/helpers/optimizers.jl

@@ -2,8 +2,6 @@
 # We can remove this once https://github.com/FluxML/Optimisers.jl/issues/205 is resolved.
 module ReactantCompatibleOptimisers
 
-using ConcreteStructs: @concrete
-using Functors: fmap
 using Optimisers: Optimisers, AbstractRule
 
 using ..Lux: Lux, Utils
@@ -21,32 +19,16 @@ function make_reactant_compatible(leaf::Optimisers.Leaf{<:AbstractRule})
     return Optimisers.Leaf(rule, state, leaf.frozen)
 end
 
-function make_reactant_compatible(opt::Optimisers.OptimiserChain, state)
-    res = make_reactant_compatible.(opt.opts, state)
-    new_opts = first.(res)
-    new_state = last.(res)
-    return Optimisers.OptimiserChain(new_opts...), new_state
+function make_reactant_compatible(opt::Optimisers.OptimiserChain)
+    return Optimisers.OptimiserChain(make_reactant_compatible.(opt.opts)...)
 end
 
-function make_reactant_compatible(opt::Optimisers.AbstractRule, state)
-    return (
-        Utils.to_rarray(opt; track_numbers = AbstractFloat),
-        Utils.to_rarray(state; track_numbers = AbstractFloat),
-    )
+function make_reactant_compatible(opt::Optimisers.AbstractRule)
+    return Utils.to_rarray(opt; track_numbers = AbstractFloat)
 end
 
 function make_reactant_compatible(opt::Optimisers.AccumGrad, state)
-    return (
-        AccumGrad(Utils.to_rarray(opt.n; track_numbers = Integer)),
-        Utils.to_rarray(state; track_numbers = Integer),
-    )
+    return Utils.to_rarray(opt.n; track_numbers = Integer)
 end
 
-@concrete struct AccumGrad <: AbstractRule
-    n
-end
-
-# XXX: the counter needs to match the client / device?
-Optimisers.init(::AccumGrad, x) = zero(x), Utils.to_rarray(1; track_numbers = Integer)
-
 end

src/helpers/training.jl

@@ -63,17 +63,17 @@ Constructor for [`TrainState`](@ref).
 [`TrainState`](@ref) object.
 """
 function TrainState(model::AbstractLuxLayer, ps, st, optimizer::Optimisers.AbstractRule)
-    st_opt = Optimisers.setup(optimizer, ps)
     if get_device_type(ps) <: ReactantDevice
-        st_opt = fmap(
-            ReactantCompatibleOptimisers.make_reactant_compatible,
-            st_opt;
-            exclude = Base.Fix2(isa, Optimisers.Leaf)
-        )
+        optimizer = ReactantCompatibleOptimisers.make_reactant_compatible(optimizer)
+        st_opt = optimisers_setup_with_jit(optimizer, ps)
+    else
+        st_opt = Optimisers.setup(optimizer, ps)
     end
     return TrainState(nothing, nothing, model, ps, st, optimizer, st_opt, 0)
 end
 
+function optimisers_setup_with_jit end
+
 @concrete struct TrainingBackendCache
     backend
     first_try <: StaticBool