Fix method ambiguity and skip some test for enzyme

Author: darsnack

src/deprecations.jl

@@ -107,7 +107,10 @@ train!(loss, ps::Params, data, opt::Optimisers.AbstractRule; cb=nothing) = error
   But better to use the new explicit style, in which `m` itself is the 2nd argument.
   """)
 
-train!(loss, model, data, opt::Optimise.AbstractOptimiser; cb=nothing) = train!(loss, model, data, _old_to_new(opt); cb)
+train!(loss, model, data, opt::Optimise.AbstractOptimiser; cb=nothing) =
+  train!(loss, model, data, _old_to_new(opt); cb)
+train!(loss, model::Enzyme.Duplicated, data, opt::Optimise.AbstractOptimiser; cb=nothing) =
+  train!(loss, model, data, _old_to_new(opt); cb)
 
 # Next, to use the new `setup` with the still-exported old-style `Adam` etc:
 import .Train: setup

src/train.jl

@@ -110,7 +110,7 @@ function train!(loss, model, data, opt; cb = nothing)
                             For more control use a loop with `gradient` and `update!`.""")
   @withprogress for (i,d) in enumerate(data)
     d_splat = d isa Tuple ? d : (d,)
-    
+
     l, gs = Zygote.withgradient(m -> loss(m, d_splat...), model)
 
     if !isfinite(l)
@@ -127,7 +127,7 @@ function train!(loss, model::Enzyme.Duplicated, data, opt; cb = nothing)
                             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)...)
 

test/train.jl

@@ -84,15 +84,15 @@ end
   y1 = m(x)  # before
 
   # Implicit gradient
-  gold = gradient(() -> m(x), Flux.params(m))
+  gold = Zygote.gradient(() -> m(x), Flux.params(m))
   @test gold isa Flux.Zygote.Grads
   @test_throws ErrorException Flux.update!(Flux.Adam(), m, gold)  # friendly
   Flux.update!(Flux.Adam(), Flux.params(m), gold)
   y2 = m(x)
   @test y2 < y1
 
   # Explicit gradient
-  gs = gradient(marg -> marg(x), m)
+  gs = Zygote.gradient(marg -> marg(x), m)
   @test gs isa Tuple
   @test_throws ErrorException Flux.update!(Flux.Adam(), Flux.params(m), gs) # friendly
   @test_throws ErrorException Flux.update!(Flux.Adam(), Flux.params(m), gs[1]) # friendly
@@ -133,17 +133,20 @@ for (trainfn!, name) in ((Flux.train!, "Zygote"), (train_enzyme!, "Enzyme"))
   diff1 = model.weight .- init_weight
 
   # Take 2: the same, but with Flux.params. Was broken for a bit, no tests!
-  model.weight .= init_weight
-  model.bias .= 0
-  pen2(x::AbstractArray) = sum(abs2, x)/2
-  opt = Flux.setup(Adam(0.1), model)
-  trainfn!(model, data, opt) do m, x, y
-    err = Flux.mse(m(x), y)
-    l2 = sum(pen2, Flux.params(m))
-    err + 0.33 * l2
+  # skipping this test for Enzyme cause implicit params is unsupported
+  if name == "Zygote"
+    model.weight .= init_weight
+    model.bias .= 0
+    pen2(x::AbstractArray) = sum(abs2, x)/2
+    opt = Flux.setup(Adam(0.1), model)
+    trainfn!(model, data, opt) do m, x, y
+      err = Flux.mse(m(x), y)
+      l2 = sum(pen2, Flux.params(m))
+      err + 0.33 * l2
+    end
+    diff2 = model.weight .- init_weight
+    @test diff1 ≈ diff2
   end
-  diff2 = model.weight .- init_weight
-  @test diff1 ≈ diff2
 
   # Take 3: using WeightDecay instead. Need the /2 above, to match exactly.
   model.weight .= init_weight