teach Flux.state about Duplicated

src/loading.jl

@@ -176,3 +176,14 @@ const STATE_TYPES = Union{AbstractArray, Number, Nothing, AbstractString, Symbol
 
 _state(x::STATE_TYPES) = x
 _state(x) = ()
+
+#=
+Starting with `gradient(f, m) == gradient(f, Duplicated(m))`,
+we choose to regard `Duplicated` as some kind of label, not part of the model tree,
+and avoid outer NamedTuples like `(; val=..., dval=...)`.
+We certainly don't want to save model gradients alongside parameters/settings:
+=#
+state(x::EnzymeCore.Duplicated) = state(x.val)
+
+loadmodel!(dst::EnzymeCore.Duplicated, src::EnzymeCore.Duplicated; kw...) = @invoke loadmodel!(dst::Any, src::Any; kw...)
+loadmodel!(dst::EnzymeCore.Duplicated, src; kw...) = (loadmodel!(dst.val, src; kw...); dst)

test/ext_enzyme/enzyme.jl

@@ -185,6 +185,15 @@ end
     # setup understands Duplicated:
     @test Flux.setup(Adam(), m1) == Flux.setup(Adam(), m1.val)
 
+    # state, loadmodel do too -- all ignore the dval branch, no outer (; val, dval) namedtuple
+    @test Flux.state(m1) == Flux.state(m1.val)
+    oldmodel = deepcopy(m1)
+    oldpar = deepcopy(Flux.state(m1))
+    m1.val.weight .= 0
+    @test Flux.loadmodel!(m1, oldmodel).val.weight ≈ oldpar.weight
+    m1.val.weight .= 0
+    @test Flux.loadmodel!(m1, oldpar).val.weight ≈ oldpar.weight
+
     # At least one Duplicated is required:
     @test_throws ArgumentError Flux.gradient(m -> sum(m.bias), Const(m1.val))
     @test_throws ArgumentError Flux.gradient((m,x) -> sum(m(x)), Const(m1.val), [1,2,3f0])