Add jacobian!! and derivative!! for mutating functions (#160)

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_allocating.jl

@@ -73,24 +73,20 @@ end
 
 DI.prepare_gradient(f, ::AutoReverseEnzyme, x) = NoGradientExtras()
 
-function DI.gradient(f, ::AutoReverseEnzyme, x::AbstractArray, ::NoGradientExtras)
+function DI.gradient(f, ::AutoReverseEnzyme, x, ::NoGradientExtras)
     return gradient(Reverse, f, x)
 end
 
-function DI.gradient!!(f, grad, ::AutoReverseEnzyme, x::AbstractArray, ::NoGradientExtras)
+function DI.gradient!!(f, grad, ::AutoReverseEnzyme, x, ::NoGradientExtras)
     grad_sametype = convert(typeof(x), grad)
     gradient!(Reverse, grad_sametype, f, x)
     return grad_sametype
 end
 
-function DI.value_and_gradient(
-    f, backend::AutoReverseEnzyme, x::AbstractArray, ::NoGradientExtras
-)
+function DI.value_and_gradient(f, backend::AutoReverseEnzyme, x, ::NoGradientExtras)
     return DI.value_and_pullback(f, backend, x, one(eltype(x)), NoPullbackExtras())
 end
 
-function DI.value_and_gradient!!(
-    f, grad, backend::AutoReverseEnzyme, x::AbstractArray, ::NoGradientExtras
-)
+function DI.value_and_gradient!!(f, grad, backend::AutoReverseEnzyme, x, ::NoGradientExtras)
     return DI.value_and_pullback!!(f, grad, backend, x, one(eltype(x)), NoPullbackExtras())
 end

DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/mutating.jl

@@ -35,6 +35,20 @@ function DI.value_and_pushforward!!(
     return y, dy
 end
 
+function DI.pushforward!!(
+    f!,
+    y,
+    dy,
+    ::AnyAutoFastDifferentiation,
+    x,
+    dx,
+    extras::FastDifferentiationMutatingPushforwardExtras,
+)
+    v_vec = vcat(myvec(x), myvec(dx))
+    extras.jvp_exe!(vec(dy), v_vec)
+    return dy
+end
+
 ## Derivative
 
 struct FastDifferentiationMutatingDerivativeExtras{E} <: DerivativeExtras
@@ -66,6 +80,18 @@ function DI.value_and_derivative!!(
     return y, der
 end
 
+function DI.derivative!!(
+    f!,
+    y,
+    der,
+    ::AnyAutoFastDifferentiation,
+    x,
+    extras::FastDifferentiationMutatingDerivativeExtras,
+)
+    extras.der_exe!(der, monovec(x))
+    return der
+end
+
 ## Jacobian
 
 struct FastDifferentiationMutatingJacobianExtras{E} <: JacobianExtras
@@ -100,3 +126,15 @@ function DI.value_and_jacobian!!(
     extras.jac_exe!(jac, vec(x))
     return y, jac
 end
+
+function DI.jacobian!!(
+    f!,
+    y,
+    jac,
+    ::AnyAutoFastDifferentiation,
+    x,
+    extras::FastDifferentiationMutatingJacobianExtras,
+)
+    extras.jac_exe!(jac, vec(x))
+    return jac
+end

DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/mutating.jl

@@ -47,6 +47,18 @@ function DI.value_and_derivative!!(
     return y, der
 end
 
+function DI.derivative!!(
+    f!,
+    y::AbstractArray,
+    der::AbstractArray,
+    backend::AnyAutoFiniteDiff,
+    x,
+    ::FiniteDiffMutatingDerivativeExtras,
+)
+    finite_difference_gradient!(der, f!, x, fdtype(backend), eltype(y), FUNCTION_INPLACE)
+    return der
+end
+
 ## Jacobian
 
 struct FiniteDiffMutatingJacobianExtras{C}
@@ -73,3 +85,15 @@ function DI.value_and_jacobian!!(
     f!(y, x)
     return y, jac
 end
+
+function DI.jacobian!!(
+    f!,
+    y::AbstractArray,
+    jac::AbstractMatrix,
+    ::AnyAutoFiniteDiff,
+    x,
+    extras::FiniteDiffMutatingJacobianExtras,
+)
+    finite_difference_jacobian!(jac, f!, x, extras.cache)
+    return jac
+end

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/mutating.jl

@@ -35,6 +35,18 @@ function DI.value_and_derivative!!(
     return DiffResults.value(result), DiffResults.derivative(result)
 end
 
+function DI.derivative!!(
+    f!,
+    y::AbstractArray,
+    der::AbstractArray,
+    ::AnyAutoForwardDiff,
+    x::Number,
+    extras::ForwardDiffMutatingDerivativeExtras,
+)
+    der = derivative!(der, f!, y, x, extras.config)
+    return der
+end
+
 ## Jacobian
 
 struct ForwardDiffMutatingJacobianExtras{C} <: JacobianExtras
@@ -61,3 +73,15 @@ function DI.value_and_jacobian!!(
     result = jacobian!(result, f!, y, x, extras.config)
     return DiffResults.value(result), DiffResults.jacobian(result)
 end
+
+function DI.jacobian!!(
+    f!,
+    y::AbstractArray,
+    jac::AbstractMatrix,
+    ::AnyAutoForwardDiff,
+    x::AbstractArray,
+    extras::ForwardDiffMutatingJacobianExtras,
+)
+    jac = jacobian!(jac, f!, y, x, extras.config)
+    return jac
+end

DifferentiationInterface/ext/DifferentiationInterfacePolyesterForwardDiffExt/mutating.jl

@@ -22,6 +22,12 @@ function DI.value_and_derivative!!(
     return DI.value_and_derivative!!(f!, y, der, single_threaded(backend), x, extras)
 end
 
+function DI.derivative!!(
+    f!, y, der, backend::AnyAutoPolyForwardDiff, x, extras::DerivativeExtras
+)
+    return DI.derivative!!(f!, y, der, single_threaded(backend), x, extras)
+end
+
 ## Jacobian
 
 DI.prepare_jacobian(f!, ::AnyAutoPolyForwardDiff, y, x) = NoJacobianExtras()
@@ -34,7 +40,19 @@ function DI.value_and_jacobian!!(
     x::AbstractArray,
     ::NoJacobianExtras,
 ) where {C}
-    f!(y, x)
     threaded_jacobian!(f!, y, jac, x, Chunk{C}())
+    f!(y, x)
     return y, jac
 end
+
+function DI.jacobian!!(
+    f!,
+    y::AbstractArray,
+    jac::AbstractMatrix,
+    ::AnyAutoPolyForwardDiff{C},
+    x::AbstractArray,
+    ::NoJacobianExtras,
+) where {C}
+    threaded_jacobian!(f!, y, jac, x, Chunk{C}())
+    return jac
+end

DifferentiationInterface/ext/DifferentiationInterfaceReverseDiffExt/mutating.jl

@@ -23,6 +23,24 @@ function DI.value_and_pullback!!(
     return y, dx
 end
 
+function DI.pullback!!(
+    f!,
+    y::AbstractArray,
+    dx::AbstractArray,
+    ::AnyAutoReverseDiff,
+    x::AbstractArray,
+    dy::AbstractArray,
+    ::NoPullbackExtras,
+)
+    function dotproduct_closure(x)
+        y_copy = similar(y, eltype(x))
+        f!(y_copy, x)
+        return dot(y_copy, dy)
+    end
+    dx = gradient!(dx, dotproduct_closure, x)
+    return dx
+end
+
 ### Number in, not supported
 
 function DI.value_and_pullback!!(
@@ -72,3 +90,15 @@ function DI.value_and_jacobian!!(
     result = jacobian!(result, extras.tape, x)
     return DiffResults.value(result), DiffResults.derivative(result)
 end
+
+function DI.jacobian!!(
+    _f!,
+    y::AbstractArray,
+    jac::AbstractMatrix,
+    ::AnyAutoReverseDiff,
+    x::AbstractArray,
+    extras::ReverseDiffMutatingJacobianExtras,
+)
+    jac = jacobian!(jac, extras.tape, x)
+    return jac
+end

DifferentiationInterface/ext/DifferentiationInterfaceSparseDiffToolsExt/mutating.jl

@@ -27,5 +27,17 @@ for AutoSparse in SPARSE_BACKENDS
             f!(y, x)
             return y, jac
         end
+
+        function DI.jacobian!!(
+            f!,
+            y,
+            jac,
+            backend::$AutoSparse,
+            x,
+            extras::SparseDiffToolsMutatingJacobianExtras,
+        )
+            sparse_jacobian!(jac, backend, extras.cache, f!, y, x)
+            return jac
+        end
     end
 end

DifferentiationInterface/src/derivative.jl

@@ -96,3 +96,17 @@ function value_and_derivative!!(
         f!, y, der, backend, x, one(x), extras.pushforward_extras
     )
 end
+
+"""
+    derivative!!(f!, y, der, backend, x, [extras]) -> der
+"""
+function derivative!!(
+    f!,
+    y,
+    der,
+    backend::AbstractADType,
+    x,
+    extras::DerivativeExtras=prepare_derivative(f!, backend, y, x),
+)
+    return pushforward!!(f!, y, der, backend, x, one(x), extras.pushforward_extras)
+end

DifferentiationInterface/src/jacobian.jl

@@ -203,3 +203,17 @@ function value_and_jacobian_aux!!(
     end
     return y, jac
 end
+
+"""
+    jacobian!!(f!, y, jac, backend, x, [extras]) -> jac
+"""
+function jacobian!!(
+    f!,
+    y,
+    jac,
+    backend::AbstractADType,
+    x,
+    extras::JacobianExtras=prepare_jacobian(f!, backend, y, x),
+)
+    return value_and_jacobian!!(f!, y, jac, backend, x, extras)[2]
+end

DifferentiationInterfaceTest/src/tests/benchmark.jl

@@ -245,15 +245,19 @@ function run_benchmark!(
     bench1 = @be (mysimilar(y), mysimilar(y)) value_and_derivative!!(
         f!, _[1], _[2], ba, x, extras
     )
+    bench2 = @be (mysimilar(y), mysimilar(y)) derivative!!(f!, _[1], _[2], ba, x, extras)
     # count
     cc! = CallCounter(f!)
     extras = prepare_derivative(cc!, ba, y, x)
     calls0 = reset_count!(cc!)
     value_and_derivative!!(cc!, mysimilar(y), mysimilar(y), ba, x, extras)
     calls1 = reset_count!(cc!)
+    derivative!!(cc!, mysimilar(y), mysimilar(y), ba, x, extras)
+    calls2 = reset_count!(cc!)
     # record
     record!(data, ba, scen, prepare_derivative, bench0, calls0)
     record!(data, ba, scen, value_and_derivative!!, bench1, calls1)
+    record!(data, ba, scen, derivative!!, bench2, calls2)
     return nothing
 end
 
@@ -322,15 +326,21 @@ function run_benchmark!(
     bench1 = @be (mysimilar(y), mysimilar(jac_template)) value_and_jacobian!!(
         f!, _[1], _[2], ba, x, extras
     )
+    bench2 = @be (mysimilar(y), mysimilar(jac_template)) jacobian!!(
+        f!, _[1], _[2], ba, x, extras
+    )
     # count
     cc! = CallCounter(f!)
     extras = prepare_jacobian(cc!, ba, y, x)
     calls0 = reset_count!(cc!)
     value_and_jacobian!!(cc!, mysimilar(y), mysimilar(jac_template), ba, x, extras)
     calls1 = reset_count!(cc!)
+    jacobian!!(cc!, mysimilar(y), mysimilar(jac_template), ba, x, extras)
+    calls2 = reset_count!(cc!)
     # record
     record!(data, ba, scen, prepare_jacobian, bench0, calls0)
     record!(data, ba, scen, value_and_jacobian!!, bench1, calls1)
+    record!(data, ba, scen, jacobian!!, bench2, calls2)
     return nothing
 end