fix: clipnorm support for reactant (#1469)

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "1.21.1"
+version = "1.21.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/helpers/optimizers.jl

@@ -70,6 +70,18 @@ function make_reactant_compatible(
     )
 end
 
+function make_reactant_compatible(
+    opt::Optimisers.ClipNorm, dev::ReactantDevice, outermost=Val(true)
+)
+    opt_ra = Optimisers.ClipNorm(
+        Utils.to_rarray(opt.omega; track_numbers=Integer, _dev_to_kwargs(dev)...),
+        opt.p,
+        false,
+    )
+    outermost isa Val{true} && return ReactantOptimiser(opt_ra)
+    return opt_ra
+end
+
 function optimisers_setup_with_jit end
 
 end

test/reactant/training_tests.jl

@@ -106,3 +106,25 @@ end
     hlo = repr(@code_hlo(Optimisers.update(st_opt, ps, ps)))
     @test length(findall("stablehlo.if", hlo)) == (2 + 1 + 2) * 2
 end
+
+@testitem "Reactant Optimisers Patch: ClipNorm" tags = [:reactant] setup = [SharedTestSetup] begin
+    using Lux, Random, Reactant, Optimisers
+
+    dev = reactant_device(; force=true)
+
+    model = Chain(
+        Dense(2 => 4, relu), Chain(Dense(4 => 2, relu; use_bias=false), Dense(2 => 2))
+    )
+    ps, st = Lux.setup(Random.default_rng(), model) |> dev
+
+    x = randn(Float32, 2, 32) |> dev
+
+    train_state = Training.TrainState(
+        model, ps, st, OptimiserChain(ClipNorm(0.5), Descent(0.1))
+    )
+
+    _, loss, stats, ts = Training.single_train_step(
+        AutoEnzyme(), MSELoss(), (x, x), train_state; return_gradients=Val(false)
+    )
+    @test loss isa Number
+end