Improve Enzyme extension, add hessians to ForwardDiff and ReverseDiff (#118)

* Improve Enzyme extension, add hessians to ForwardDiff and ReverseDiff

* Remove Enzyme forward from type stable backends

Author: gdalle

docs/src/overview.md

@@ -36,15 +36,19 @@ Several variants of each operator are defined:
 !!! warning
     The "bang-bang" syntactic convention `!!` signals that some of the arguments _can_ be mutated, but they do not _have to be_.
     Such arguments will always be part of the return, so that one can simply reuse the operator's output and forget its input.
-
     In other words, this is good:
     ```julia
-    grad = gradient!!(f, grad, backend, x)  # do this
+    # work with grad_in
+    grad_out = gradient!!(f, grad_in, backend, x)
+    # work with grad_out
     ```
-    On the other hand, this is bad, because if `grad` has not been mutated, you will get wrong results:
+    On the other hand, this is bad, because if `grad_in` has not been mutated, you will forget the results:
     ```julia
-    gradient!!(f, grad, backend, x)  # don't do this
+    # work with grad_in
+    gradient!!(f, grad_in, backend, x)
+    # mistakenly keep working with grad_in
     ```
+    Note that we don't guarantee `grad_out` will have the same type as `grad_in`.
 
 ## Second order
 

ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl

@@ -1,7 +1,6 @@
 module DifferentiationInterfaceEnzymeExt
 
 using ADTypes: ADTypes, AutoEnzyme
-using DifferentiationInterface: myupdate!!
 import DifferentiationInterface as DI
 using DocStringExtensions
 using Enzyme:
@@ -49,6 +48,12 @@ function DI.basis(::AutoEnzyme, a::AbstractArray{T}, i::CartesianIndex) where {T
     return b
 end
 
+function zero_sametype!!(x_target, x)
+    x_sametype = convert(typeof(x), x_target)
+    x_sametype .= zero(eltype(x))
+    return x_sametype
+end
+
 include("forward_allocating.jl")
 include("forward_mutating.jl")
 

ext/DifferentiationInterfaceEnzymeExt/forward_allocating.jl

@@ -1,7 +1,76 @@
 ## Pushforward
 
 function DI.value_and_pushforward(f, backend::AutoForwardEnzyme, x, dx, extras::Nothing)
-    dx_sametype = convert(typeof(x), copy(dx))
+    dx_sametype = convert(typeof(x), dx)
     y, new_dy = autodiff(backend.mode, f, Duplicated, Duplicated(x, dx_sametype))
     return y, new_dy
 end
+
+function DI.pushforward(f, backend::AutoForwardEnzyme, x, dx, extras::Nothing)
+    dx_sametype = convert(typeof(x), dx)
+    new_dy = only(autodiff(backend.mode, f, DuplicatedNoNeed, Duplicated(x, dx_sametype)))
+    return new_dy
+end
+
+function DI.value_and_pushforward!!(
+    f, _dy, backend::AutoForwardEnzyme, x, dx, extras::Nothing
+)
+    # dy cannot be passed anyway
+    return DI.value_and_pushforward(f, backend, x, dx, extras)
+end
+
+function DI.pushforward!!(f, _dy, backend::AutoForwardEnzyme, x, dx, extras::Nothing)
+    # dy cannot be passed anyway
+    return DI.pushforward(f, backend, x, dx, extras)
+end
+
+## Gradient
+
+function DI.gradient(f, backend::AutoForwardEnzyme, x::AbstractArray, extras::Nothing)
+    return reshape(collect(gradient(backend.mode, f, x)), size(x))
+end
+
+function DI.value_and_gradient(
+    f, backend::AutoForwardEnzyme, x::AbstractArray, extras::Nothing
+)
+    return f(x), DI.gradient(f, backend, x, extras)
+end
+
+function DI.gradient!!(
+    f, _grad, backend::AutoForwardEnzyme, x::AbstractArray, extras::Nothing
+)
+    return DI.gradient(f, backend, x, extras)
+end
+
+function DI.value_and_gradient!!(
+    f, _grad, backend::AutoForwardEnzyme, x::AbstractArray, extras::Nothing
+)
+    return DI.value_and_gradient(f, backend, x, extras)
+end
+
+## Jacobian
+
+function DI.jacobian(f, backend::AutoForwardEnzyme, x::AbstractArray, extras::Nothing)
+    jac_wrongshape = jacobian(backend.mode, f, x)
+    nx = length(x)
+    ny = length(jac_wrongshape) ÷ length(x)
+    return reshape(jac_wrongshape, ny, nx)
+end
+
+function DI.value_and_jacobian(
+    f, backend::AutoForwardEnzyme, x::AbstractArray, extras::Nothing
+)
+    return f(x), DI.jacobian(f, backend, x, extras)
+end
+
+function DI.jacobian!!(
+    f, _jac, backend::AutoForwardEnzyme, x::AbstractArray, extras::Nothing
+)
+    return DI.jacobian(f, backend, x, extras)
+end
+
+function DI.value_and_jacobian!!(
+    f, _jac, backend::AutoForwardEnzyme, x::AbstractArray, extras::Nothing
+)
+    return DI.value_and_jacobian(f, backend, x, extras)
+end

ext/DifferentiationInterfaceEnzymeExt/forward_mutating.jl

@@ -3,10 +3,10 @@
 function DI.value_and_pushforward!!(
     f!, y, dy, backend::AutoForwardEnzyme, x, dx, extras::Nothing
 )
-    dx_sametype = convert(typeof(x), copy(dx))
-    dy_sametype = convert(typeof(y), dy)
+    dx_sametype = convert(typeof(x), dx)
+    dy_sametype = zero_sametype!!(dy, y)
     autodiff(
         backend.mode, f!, Const, Duplicated(y, dy_sametype), Duplicated(x, dx_sametype)
     )
-    return y, myupdate!!(dy, dy_sametype)
+    return y, dy_sametype
 end

ext/DifferentiationInterfaceEnzymeExt/reverse_allocating.jl

@@ -1,5 +1,7 @@
 ## Pullback
 
+### Out-of-place
+
 function DI.value_and_pullback(
     f, ::AutoReverseEnzyme, x::Number, dy::Number, extras::Nothing
 )
@@ -20,28 +22,28 @@ function DI.value_and_pullback(
     return y, new_dx
 end
 
+### In-place
+
 function DI.value_and_pullback!!(
     f, dx, ::AutoReverseEnzyme, x::AbstractArray, dy::Number, extras::Nothing
 )
-    dx_sametype = convert(typeof(x), dx)
-    dx_sametype .= zero(eltype(dx_sametype))
+    dx_sametype = zero_sametype!!(dx, x)
     _, y = autodiff(ReverseWithPrimal, f, Active, Duplicated(x, dx_sametype))
     dx_sametype .*= dy
-    return y, myupdate!!(dx, dx_sametype)
+    return y, dx_sametype
 end
 
 function DI.value_and_pullback!!(
     f, dx, ::AutoReverseEnzyme, x::AbstractArray, dy::AbstractArray, extras::Nothing
 )
-    dx_sametype = convert(typeof(x), dx)
-    dx_sametype .= zero(eltype(dx_sametype))
+    dx_sametype = zero_sametype!!(dx, x)
     forw, rev = autodiff_thunk(
         ReverseSplitWithPrimal, Const{typeof(f)}, Duplicated, Duplicated{typeof(x)}
     )
     tape, y, new_dy = forw(Const(f), Duplicated(x, dx_sametype))
     new_dy .= dy
     rev(Const(f), Duplicated(x, dx_sametype), tape)
-    return y, myupdate!!(dx, dx_sametype)
+    return y, dx_sametype
 end
 
 function DI.value_and_pullback(f, backend::AutoReverseEnzyme, x::AbstractArray, dy, extras)
@@ -58,6 +60,5 @@ end
 function DI.gradient!!(f, grad, ::AutoReverseEnzyme, x::AbstractArray, extras::Nothing)
     grad_sametype = convert(typeof(x), grad)
     gradient!(Reverse, grad_sametype, f, x)
-    grad .= grad_sametype
-    return grad
+    return grad_sametype
 end

ext/DifferentiationInterfaceEnzymeExt/reverse_mutating.jl

@@ -11,9 +11,8 @@ end
 function DI.value_and_pullback!!(
     f!, y, dx, ::AutoReverseEnzyme, x::AbstractArray, dy, extras::Nothing
 )
-    dx_sametype = convert(typeof(x), dx)
-    dx_sametype .= zero(eltype(dx_sametype))
-    dy_sametype = convert(typeof(y), copy(dy))  # TODO: how to get rid of copy?
+    dx_sametype = zero_sametype!!(dx, x)
+    dy_sametype = convert(typeof(y), copy(dy))
     autodiff(Reverse, f!, Const, Duplicated(y, dy_sametype), Duplicated(x, dx_sametype))
-    return y, myupdate!!(dx, dx_sametype)
+    return y, dx_sametype
 end

ext/DifferentiationInterfaceForwardDiffExt/DifferentiationInterfaceForwardDiffExt.jl

@@ -9,6 +9,7 @@ using ForwardDiff:
     DerivativeConfig,
     ForwardDiff,
     GradientConfig,
+    HessianConfig,
     JacobianConfig,
     Tag,
     derivative,
@@ -17,6 +18,8 @@ using ForwardDiff:
     extract_derivative!,
     gradient,
     gradient!,
+    hessian,
+    hessian!,
     jacobian,
     jacobian!,
     value

ext/DifferentiationInterfaceForwardDiffExt/allocating.jl

@@ -70,3 +70,19 @@ end
 function DI.jacobian(f, ::AutoForwardDiff, x::AbstractArray, config::JacobianConfig)
     return jacobian(f, x, config)
 end
+
+## Hessian
+
+function DI.prepare_hessian(f, backend::AutoForwardDiff, x::AbstractArray)
+    return HessianConfig(f, x, choose_chunk(backend, x))
+end
+
+function DI.hessian!!(
+    f, hess::AbstractMatrix, ::AutoForwardDiff, x::AbstractArray, config::HessianConfig
+)
+    return hessian!(hess, f, x, config)
+end
+
+function DI.hessian(f, ::AutoForwardDiff, x::AbstractArray, config::HessianConfig)
+    return hessian(f, x, config)
+end

ext/DifferentiationInterfaceReverseDiffExt/DifferentiationInterfaceReverseDiffExt.jl

@@ -7,14 +7,19 @@ using DocStringExtensions
 using LinearAlgebra: mul!
 using ReverseDiff:
     CompiledGradient,
+    CompiledHessian,
     CompiledJacobian,
     GradientConfig,
     GradientTape,
+    HessianConfig,
+    HessianTape,
     JacobianConfig,
     JacobianTape,
     compile,
     gradient,
     gradient!,
+    hessian,
+    hessian!,
     jacobian,
     jacobian!
 

ext/DifferentiationInterfaceReverseDiffExt/allocating.jl

@@ -141,3 +141,29 @@ function DI.jacobian(
 )
     return jacobian!(tape, x)
 end
+
+## Hessian
+
+function DI.prepare_hessian(f, backend::AutoReverseDiff, x::AbstractArray)
+    tape = HessianTape(f, x)
+    if backend.compile
+        tape = compile(tape)
+    end
+    return tape
+end
+
+function DI.hessian!!(
+    _f,
+    hess::AbstractMatrix,
+    ::AutoReverseDiff,
+    x::AbstractArray,
+    tape::Union{HessianTape,CompiledHessian},
+)
+    return hessian!(hess, tape, x)
+end
+
+function DI.hessian(
+    _f, ::AutoReverseDiff, x::AbstractArray, tape::Union{HessianTape,CompiledHessian}
+)
+    return hessian!(tape, x)
+end