Use Mooncake value_and_gradient!! helpers

Author: yebai

Project.toml

@@ -63,7 +63,7 @@ IrrationalConstants = "0.1, 0.2"
 LazyArrays = "2"
 LogExpFunctions = "0.3.3"
 MappedArrays = "0.2.2, 0.3, 0.4"
-Mooncake = "0.4.95, 0.5"
+Mooncake = "0.5.26"
 Reexport = "0.2, 1"
 ReverseDiff = "1"
 Roots = "1.3.15, 2"

ext/BijectorsMooncakeExt.jl

@@ -68,28 +68,10 @@ end
 
 ## Forward-mode implementations (column-by-column JVPs)
 
-function _value_and_gradient(
-    f, backend::AutoMooncakeForward, x::AbstractVector{T}
-) where {T}
-    val = f(x)
-    n = length(x)
-    cache = prepare_derivative_cache(f, x; config=_mooncake_config(backend))
-    df = _mooncake_zero_tangent_or_primal(f, backend)
-    if n == 0
-        return val, Vector{eltype(x)}(undef, 0)
-    end
-    dx = zeros(T, n)
-    dx[1] = one(T)
-    _, first_jvp = value_and_derivative!!(cache, (f, df), (x, dx))
-    grad = Vector{typeof(first_jvp)}(undef, n)
-    grad[1] = first_jvp
-    for j in 2:n
-        fill!(dx, zero(T))
-        dx[j] = one(T)
-        _, jvp = value_and_derivative!!(cache, (f, df), (x, dx))
-        grad[j] = jvp
-    end
-    return val, grad
+function _value_and_gradient(f, backend::AutoMooncakeForward, x::AbstractVector)
+    cache = prepare_gradient_cache(f, x; config=_mooncake_config(backend))
+    val, (_, x_grad) = value_and_gradient!!(cache, f, x)
+    return val, x_grad
 end
 
 function _value_and_jacobian(

test/Project.toml

@@ -53,7 +53,7 @@ LazyArrays = "1, 2"
 LogDensityProblems = "2"
 LogExpFunctions = "0.3.1"
 MCMCDiagnosticTools = "0.3"
-Mooncake = "0.4, 0.5"
+Mooncake = "0.5.26"
 PDMats = "0.11"
 ReverseDiff = "1.4.2"
 StableRNGs = "1"