Add AD.value_and_gradient too

FluxML · Oct 15, 2022 · 3fb5410 · 3fb5410
1 parent 8b8eebc
commit 3fb5410
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/src/optimise/gradients.jl b/src/optimise/gradients.jl
@@ -9,6 +9,8 @@ AD.@primitive pullback_function(ad::ZygoteImplicitBackend, f, x::Zygote.Params)
 # this is a hack to get around
 # https://github.com/JuliaDiff/AbstractDifferentiation.jl/issues/63#issuecomment-1225959150
 AD.gradient(::ZygoteImplicitBackend, f, x::Zygote.Params) = Zygote.gradient(f, x)
+AD.value_and_gradient(::ZygoteImplicitBackend, f, x::Zygote.Params) =
+    Zygote.withgradient(f, x)
 
 struct ZygoteExplicitBackend{T} <: AD.AbstractReverseMode
     core_backend::T
@@ -21,3 +23,5 @@ AD.@primitive pullback_function(ad::ZygoteExplicitBackend, f, xs...) =
 # this is a hack to get around
 # https://github.com/JuliaDiff/AbstractDifferentiation.jl/issues/63#issuecomment-1225959150
 AD.gradient(::ZygoteExplicitBackend, f, xs...) = Zygote.gradient(f, xs...)
+AD.value_and_gradient(::ZygoteExplicitBackend, f, xs...) =
+    Zygote.withgradient(f, xs...)
diff --git a/src/optimise/train.jl b/src/optimise/train.jl
@@ -94,7 +94,7 @@ _gradient_only(x) = error("Expected gradient w.r.t. single argument (or Zygote.G
 
 """
     train!(loss, pars::Params, data, opt::AbstractOptimiser; [cb])
-        
+
 Uses a `loss` function and training `data` to improve the 
 model's parameters according to a particular optimisation rule `opt`.
-Original file line number
+Diff line change
@@ Expand Up @@
     """
         train!(loss, pars::Params, data, opt::AbstractOptimiser; [cb])
     Uses a `loss` function and training `data` to improve the
     model's parameters according to a particular optimisation rule `opt`.
@@ Expand Down @@