feat: support general optimisers

avik-pal · avik-pal · commit bed8b135b190 · 2025-03-06T17:01:31.000-05:00
diff --git a/Project.toml b/Project.toml
@@ -110,7 +110,7 @@ NNlib = "0.9.26"
 Optimisers = "0.4.5"
 Preferences = "1.4.3"
 Random = "1.10"
-Reactant = "0.2.27"
+Reactant = "0.2.36"
 Reexport = "1.2.2"
 ReverseDiff = "1.15"
 SIMDTypes = "0.1"
diff --git a/docs/Project.toml b/docs/Project.toml
@@ -52,11 +52,11 @@ LuxLib = "1.3.4"
 LuxTestUtils = "1.5"
 MLDataDevices = "1.6.10"
 NNlib = "0.9.26"
-Optimisers = "0.4.1"
+Optimisers = "0.4.5"
 Pkg = "1.10"
 Printf = "1.10"
 Random = "1.10"
-Reactant = "0.2.27"
+Reactant = "0.2.36"
 StableRNGs = "1"
 StaticArrays = "1"
 WeightInitializers = "1"
diff --git a/examples/Basics/Project.toml b/examples/Basics/Project.toml
@@ -11,5 +11,5 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 ComponentArrays = "0.15.22"
 ForwardDiff = "0.10"
 Lux = "1"
-Optimisers = "0.4.1"
+Optimisers = "0.4.5"
 Zygote = "0.6.70, 0.7"
diff --git a/examples/CIFAR10/Project.toml b/examples/CIFAR10/Project.toml
@@ -35,11 +35,11 @@ LuxCUDA = "0.3.2"
 MLDatasets = "0.7.14"
 MLUtils = "0.4.4"
 OneHotArrays = "0.2.5"
-Optimisers = "0.4.1"
+Optimisers = "0.4.5"
 Printf = "1.10"
 ProgressTables = "0.1"
 Random = "1.10"
-Reactant = "0.2.27"
+Reactant = "0.2.36"
 StableRNGs = "1.0.2"
 Statistics = "1.10"
 Zygote = "0.6.70, 0.7"
diff --git a/examples/ConvolutionalVAE/Project.toml b/examples/ConvolutionalVAE/Project.toml
@@ -26,4 +26,4 @@ OneHotArrays = "0.2.6"
 Optimisers = "0.4"
 Printf = "1.10"
 Random = "1.10"
-Reactant = "0.2.27"
+Reactant = "0.2.36"
diff --git a/examples/DDIM/Project.toml b/examples/DDIM/Project.toml
@@ -36,7 +36,7 @@ JLD2 = "0.4.48, 0.5"
 Lux = "1"
 LuxCUDA = "0.3"
 MLUtils = "0.4"
-Optimisers = "0.4.1"
+Optimisers = "0.4.5"
 ParameterSchedulers = "0.4.1"
 ProgressBars = "1"
 Random = "1.10"
diff --git a/examples/GCN_Cora/Project.toml b/examples/GCN_Cora/Project.toml
@@ -18,8 +18,8 @@ GNNGraphs = "1"
 Lux = "1.5"
 MLDatasets = "0.7.18"
 OneHotArrays = "0.2"
-Optimisers = "0.4.4"
+Optimisers = "0.4.5"
 Printf = "1.10"
 Random = "1.10"
-Reactant = "0.2.27"
+Reactant = "0.2.36"
 Statistics = "1.10"
diff --git a/examples/HyperNet/Project.toml b/examples/HyperNet/Project.toml
@@ -15,5 +15,5 @@ Lux = "1"
 MLDatasets = "0.7"
 MLUtils = "0.4"
 OneHotArrays = "0.2.5"
-Optimisers = "0.4.1"
-Reactant = "0.2.27"
+Optimisers = "0.4.5"
+Reactant = "0.2.36"
diff --git a/examples/ImageNet/Project.toml b/examples/ImageNet/Project.toml
@@ -38,7 +38,7 @@ MLUtils = "0.4.4"
 MPI = "0.20.21"
 NCCL = "0.1.1"
 OneHotArrays = "0.2.5"
-Optimisers = "0.4.1"
+Optimisers = "0.4.5"
 ParameterSchedulers = "0.4.2"
 Random = "1.10"
 Setfield = "1.1.1"
diff --git a/examples/NeuralODE/Project.toml b/examples/NeuralODE/Project.toml
@@ -21,7 +21,7 @@ LuxCUDA = "0.3"
 MLDatasets = "0.7"
 MLUtils = "0.4"
 OneHotArrays = "0.2.5"
-Optimisers = "0.4.1"
+Optimisers = "0.4.5"
 OrdinaryDiffEqTsit5 = "1"
 SciMLSensitivity = "7.63"
 Statistics = "1"
diff --git a/examples/PINN2DPDE/Project.toml b/examples/PINN2DPDE/Project.toml
@@ -18,8 +18,8 @@ Enzyme = "0.13"
 Lux = "1"
 MLUtils = "0.4.4"
 OnlineStats = "1.7.1"
-Optimisers = "0.4.1"
+Optimisers = "0.4.5"
 Printf = "1.10"
 Random = "1.10"
-Reactant = "0.2.27"
+Reactant = "0.2.36"
 Statistics = "1.10"
diff --git a/examples/PolynomialFitting/Project.toml b/examples/PolynomialFitting/Project.toml
@@ -12,8 +12,8 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 ADTypes = "1.10"
 CairoMakie = "0.12, 0.13"
 Lux = "1"
-Optimisers = "0.4.1"
+Optimisers = "0.4.5"
 Printf = "1.10"
 Random = "1.10"
-Reactant = "0.2.27"
+Reactant = "0.2.36"
 Statistics = "1.10"
diff --git a/examples/RealNVP/Project.toml b/examples/RealNVP/Project.toml
@@ -16,8 +16,8 @@ ConcreteStructs = "0.2.3"
 Enzyme = "0.13.28"
 Lux = "1.5"
 MLUtils = "0.4.5"
-Optimisers = "0.4.4"
+Optimisers = "0.4.5"
 Printf = "1.10"
 Random = "1.10"
-Reactant = "0.2.27"
+Reactant = "0.2.36"
 Statistics = "1.10"
diff --git a/examples/SimpleChains/Project.toml b/examples/SimpleChains/Project.toml
@@ -17,8 +17,8 @@ Lux = "1"
 MLDatasets = "0.7.14"
 MLUtils = "0.4"
 OneHotArrays = "0.2.5"
-Optimisers = "0.4.1"
+Optimisers = "0.4.5"
 Random = "1"
-Reactant = "0.2.27"
+Reactant = "0.2.36"
 SimpleChains = "0.4.6"
 Zygote = "0.6.69, 0.7"
diff --git a/examples/SimpleRNN/Project.toml b/examples/SimpleRNN/Project.toml
@@ -13,5 +13,5 @@ ADTypes = "1.10"
 JLD2 = "0.5"
 Lux = "1"
 MLUtils = "0.4"
-Optimisers = "0.4.1"
-Reactant = "0.2.27"
+Optimisers = "0.4.5"
+Reactant = "0.2.36"
diff --git a/ext/LuxReactantExt/patches.jl b/ext/LuxReactantExt/patches.jl
@@ -2,3 +2,17 @@ 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
+    @. accum_dx += dx / opt.n
+    @trace if counter == opt.n
+        dx_final = dx
+        counter = 1
+    else
+        dx_final = zero.(dx)
+        counter += 1
+    end
+    return (accum_dx, counter), dx_final
+end
diff --git a/lib/LuxCore/Project.toml b/lib/LuxCore/Project.toml
@@ -40,7 +40,7 @@ EnzymeCore = "0.8.6"
 Functors = "0.5"
 MLDataDevices = "1.6.10"
 Random = "1.10"
-Reactant = "0.2.27"
+Reactant = "0.2.36"
 ReverseDiff = "1.15"
 Setfield = "1"
 Tracker = "0.2.36"
diff --git a/lib/LuxLib/Project.toml b/lib/LuxLib/Project.toml
@@ -92,7 +92,7 @@ Octavian = "0.3.28"
 Polyester = "0.7.15"
 Preferences = "1.4.3"
 Random = "1.10"
-Reactant = "0.2.27"
+Reactant = "0.2.36"
 Reexport = "1"
 ReverseDiff = "1.15"
 SLEEFPirates = "0.6.43"
diff --git a/lib/MLDataDevices/Project.toml b/lib/MLDataDevices/Project.toml
@@ -66,7 +66,7 @@ Metal = "1"
 OneHotArrays = "0.2.5"
 Preferences = "1.4"
 Random = "1.10"
-Reactant = "0.2.27"
+Reactant = "0.2.36"
 RecursiveArrayTools = "3.8"
 ReverseDiff = "1.15"
 SparseArrays = "1.10"
diff --git a/lib/WeightInitializers/Project.toml b/lib/WeightInitializers/Project.toml
@@ -41,7 +41,7 @@ GPUArraysCore = "0.1.6, 0.2"
 LinearAlgebra = "1.10"
 Metal = "1.3.0"
 Random = "1.10"
-Reactant = "0.2.27"
+Reactant = "0.2.36"
 SpecialFunctions = "2.4"
 Statistics = "1.10"
 julia = "1.10"
diff --git a/src/helpers/optimizers.jl b/src/helpers/optimizers.jl
@@ -4,31 +4,55 @@ module ReactantCompatibleOptimisers
 
 using ConcreteStructs: @concrete
 using Optimisers: Optimisers, AbstractRule
-using Setfield: Setfield, @set!
 
 using ..Lux: Lux, Utils
 
 abstract type ReactantCompatibleOptimisersRule <: AbstractRule end
 
 function make_reactant_compatible(opt::AbstractRule)
-    fields = fieldnames(typeof(opt))
-    for field in fields
-        opt = Setfield.set(
-            opt,
-            Setfield.PropertyLens{field}(),
-            Utils.to_rarray(getfield(opt, field); track_numbers=true)
-        )
-    end
-    return opt
+    return Utils.to_rarray(opt; track_numbers=AbstractFloat)
+end
+
+function make_reactant_compatible(opt::Optimisers.RMSProp)
+    return Optimisers.RMSProp(
+        Utils.to_rarray(opt.eta; track_numbers=AbstractFloat),
+        Utils.to_rarray(opt.rho; track_numbers=AbstractFloat),
+        opt.epsilon,
+        opt.centred
+    )
+end
+
+function make_reactant_compatible(opt::Optimisers.AdamW)
+    return Optimisers.AdamW(
+        Utils.to_rarray(opt.eta; track_numbers=AbstractFloat),
+        Utils.to_rarray(opt.beta; track_numbers=AbstractFloat),
+        Utils.to_rarray(opt.lambda; track_numbers=AbstractFloat),
+        opt.epsilon,
+        opt.coupled
+    )
+end
+
+function make_reactant_compatible(opt::Optimisers.ClipNorm)
+    return Optimisers.ClipNorm(
+        Utils.to_rarray(opt.omega; track_numbers=AbstractFloat),
+        Utils.to_rarray(opt.p; track_numbers=AbstractFloat),
+        opt.throw
+    )
 end
 
 function make_reactant_compatible(opt::Optimisers.OptimiserChain)
     return Optimisers.OptimiserChain(make_reactant_compatible.(opt.opts))
 end
 
-function make_reactant_compatible(::Optimisers.AccumGrad)
-    # The conditional will have to be traced
-    error("Reactant + AccumGrad is currently incompatible with the Lux Training API")
+function make_reactant_compatible(opt::Optimisers.AccumGrad)
+    return AccumGrad(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
diff --git a/test/Project.toml b/test/Project.toml
@@ -65,7 +65,7 @@ MLUtils = "0.4.3"
 NNlib = "0.9.26"
 Octavian = "0.3.28"
 OneHotArrays = "0.2.5"
-Optimisers = "0.4.1"
+Optimisers = "0.4.5"
 Pkg = "1.10"
 Preferences = "1.4.3"
 Random = "1.10"
