add test

ext/OptimisersEnzymeCoreExt.jl

@@ -1,13 +1,12 @@
 module OptimisersEnzymeCoreExt
 
-import Optimisers: trainable, setup, update!, isnumeric, AbstractRule
+import Optimisers: trainable, setup, update!, isnumeric, AbstractRule, _setup
 import EnzymeCore: Duplicated, Const
 
 using Functors: fmapstructure
 
-println("loaded!")
-
 trainable(x::Duplicated) = (; val = x.val)
+trainable(x::Const) = (;)
 
 """
     setup(rule::AbstractRule, model_grad::Duplicated)
@@ -16,6 +15,11 @@ For use with Enzyme's Duplicated, this just calls `setup(rule, model_grad.val)`.
 """
 setup(rule::AbstractRule, model_grad::Duplicated) = setup(rule, model_grad.val)
 
+_setup(rule, x::Duplicated; cache) = throw(ArgumentError(
+    """Objects of type `Duplicated` are only supported by Optimisers.jl at top level,
+    they may not appear deep inside other objects."""
+))
+
 """
     update!(opt_state, model_grad::Duplicated)
 

test/runtests.jl

@@ -1,5 +1,5 @@
 using Optimisers
-using ChainRulesCore, Functors, StaticArrays, Zygote
+using ChainRulesCore, Functors, StaticArrays, Zygote, EnzymeCore
 using LinearAlgebra, Statistics, Test, Random
 using Optimisers: @.., @lazy
 using Base.Broadcast: broadcasted, instantiate, Broadcasted
@@ -534,6 +534,16 @@ end
         @test Optimisers._norm(bc2, p) isa Float64
       end
     end
+
+    @testset "Enzyme Duplicated" begin
+        x_dx = Duplicated(Float16[1,2,3], Float16[1,0,-4])
+        st = Optimisers.setup(Momentum(1/9), x_dx)  # acts only on x not on dx
+        @test st isa Optimisers.Leaf
+        @test nothing === Optimisers.update!(st, x_dx)  # mutates both arguments
+        @test x_dx.val ≈ Float16[0.8887, 2.0, 3.445]
+
+        @test_throws ArgumentError Optimisers.setup(Adam(), (; a=[1,2,3.], b=x_dx))
+    end
   end
   @testset verbose=true "Destructure" begin
     include("destructure.jl")