Rearrange dispatch for enzyme train

Author: darsnack

src/functor.jl

@@ -90,45 +90,6 @@ function params!(p::Params, x, seen = IdSet())
   end
 end
 
-function Enzyme.EnzymeRules.augmented_primal(config, func::Enzyme.Const{typeof(params!)}, ::Type{RT},
-                                                        p::Enzyme.Annotation,
-                                                        x::Enzyme.Annotation,
-                                                        seen::Enzyme.Annotation) where {RT}
-
-    res = func.val(p.val, x.val, seen.val)
-
-    primal = if EnzymeRules.needs_primal(config)
-        res
-    else
-        nothing
-    end
-
-    sres = if EnzymeRules.width(config) == 1
-        func.val(p.dval, x.dval, seen isa Const ? IdSet() : seen.dval)
-    else
-        ntuple(Val(EnzymeRules.width(config))) do i
-            Base.@_inline_meta
-            func.val(p.dval[i], x.dval[i], seen isa Const ? IdSet() : seen.dval[i])
-        end
-    end
-
-    shadow = if EnzymeRules.needs_shadow(config)
-        sres
-    else
-        nothing
-    end
-
-    return EnzymeRules.AugmentedReturn(primal, shadow, nothing)
-end
-
-function Enzyme.EnzymeRules.reverse(config, func::Enzyme.Const{typeof(params!)}, ::Type{RT}, cache,
-                                      p::Enzyme.Annotation,
-                                      x::Enzyme.Annotation,
-                                      seen::Enzyme.Annotation) where {RT}
-
-    return (nothing, nothing, nothing)
-end
-
 """
     params(model)
     params(layers...)

src/train.jl

@@ -111,25 +111,32 @@ function train!(loss, model, data, opt; cb = nothing)
   @withprogress for (i,d) in enumerate(data)
     d_splat = d isa Tuple ? d : (d,)
 
-    if model isa Enzyme.Duplicated
-      _make_zero!(model.dval)
-      _, l = Enzyme.autodiff(Enzyme.ReverseWithPrimal, _applyloss, Enzyme.Active, Enzyme.Const(loss), model, map(Enzyme.Const, d_splat)...)
+    l, gs = Zygote.withgradient(m -> loss(m, d_splat...), model)
 
-      if !isfinite(l)
-        throw(DomainError(lazy"Loss is $l on data item $i, stopping training"))
-      end
-      opt, model2 = Optimisers.update!(opt, model.val, model.dval)
-      model = Enzyme.Duplicated(model2, model.dval)
-    else
-      l, gs = Zygote.withgradient(m -> loss(m, d_splat...), model)
+    if !isfinite(l)
+      throw(DomainError(lazy"Loss is $l on data item $i, stopping training"))
+    end
 
-      if !isfinite(l)
-        throw(DomainError(lazy"Loss is $l on data item $i, stopping training"))
-      end
+    opt, model = Optimisers.update!(opt, model, gs[1])
 
-      opt, model = Optimisers.update!(opt, model, gs[1])
+    @logprogress Base.haslength(data) ? i/length(data) : nothing
+  end
+end
+function train!(loss, model::Enzyme.Duplicated, data, opt; cb = nothing)
+  isnothing(cb) || error("""train! does not support callback functions.
+                            For more control use a loop with `gradient` and `update!`.""")
+  @withprogress for (i,d) in enumerate(data)
+    d_splat = d isa Tuple ? d : (d,)
+    
+    _make_zero!(model.dval)
+    _, l = Enzyme.autodiff(Enzyme.ReverseWithPrimal, _applyloss, Enzyme.Active, Enzyme.Const(loss), model, map(Enzyme.Const, d_splat)...)
 
+    if !isfinite(l)
+      throw(DomainError(lazy"Loss is $l on data item $i, stopping training"))
     end
+    opt, model2 = Optimisers.update!(opt, model.val, model.dval)
+    model = Enzyme.Duplicated(model2, model.dval)
+
     @logprogress Base.haslength(data) ? i/length(data) : nothing
   end
 end
@@ -138,6 +145,9 @@ end
 function train!(loss, model, data, rule::Optimisers.AbstractRule; cb = nothing)
   train!(loss, model, data, _rule_to_state(model, rule); cb)
 end
+function train!(loss, model::Enzyme.Duplicated, data, rule::Optimisers.AbstractRule; cb = nothing)
+  train!(loss, model, data, _rule_to_state(model, rule); cb)
+end
 
 function _rule_to_state(model, rule::Optimisers.AbstractRule)
   state = setup(rule, model)