Support second order for Enzyme

DifferentiationInterface/docs/make.jl

@@ -34,12 +34,23 @@ makedocs(;
     format=Documenter.HTML(; assets=["assets/favicon.ico"]),
     pages=[
         "Home" => "index.md",
-        "Tutorials" => ["tutorial1.md", "tutorial2.md"],
-        "Reference" => ["operators.md", "backends.md", "api.md"],
-        "Advanced" => ["preparation.md", "overloads.md"],
+        "Tutorials" => [
+            "tutorial1.md",  #
+            "tutorial2.md",
+        ],
+        "Reference" => [
+            "operators.md",  #
+            # "backends.md",
+            "api.md",
+        ],
+        "Advanced" => [
+            "preparation.md",  #
+            "overloads.md"
+        ],
     ],
     checkdocs=:exports,
     plugins=[links],
+    pagesonly=true,
 )
 
 deploydocs(;

DifferentiationInterface/docs/src/tutorial1.md

@@ -12,7 +12,13 @@ A common use case of automatic differentiation (AD) is optimizing real-valued fu
 Let's define a simple objective and a random input vector
 
 ```@example tuto1
-f(x) = sum(abs2, x)
+function f(x::AbstractVector{T}) where {T}
+    y = zero(T)
+    for i in eachindex(x)
+        y += abs2(x[i])
+    end
+    return y
+end
 
 x = collect(1.0:5.0)
 ```

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl

@@ -1,23 +1,30 @@
 module DifferentiationInterfaceEnzymeExt
 
 using ADTypes: ADTypes, AutoEnzyme
+using Compat
 import DifferentiationInterface as DI
 using DifferentiationInterface:
     DerivativeExtras,
     GradientExtras,
+    HessianExtras,
     JacobianExtras,
     PullbackExtras,
     PushforwardExtras,
+    SecondDerivativeExtras,
     NoDerivativeExtras,
     NoGradientExtras,
+    NoHessianExtras,
+    NoHVPExtras,
     NoJacobianExtras,
     NoPullbackExtras,
     NoPushforwardExtras,
+    NoSecondDerivativeExtras,
     pick_chunksize
 using DocStringExtensions
 using Enzyme:
     Active,
     Const,
+    Mode,
     Duplicated,
     DuplicatedNoNeed,
     Forward,
@@ -27,22 +34,26 @@ using Enzyme:
     ReverseSplitWithPrimal,
     ReverseMode,
     autodiff,
+    autodiff_deferred,
     autodiff_thunk,
     chunkedonehot,
     gradient,
     gradient!,
     jacobian,
     make_zero
 
+const AutoMixedEnzyme = AutoEnzyme{Nothing}
 const AutoForwardEnzyme = AutoEnzyme{<:ForwardMode}
 const AutoForwardOrNothingEnzyme = Union{AutoEnzyme{<:ForwardMode},AutoEnzyme{Nothing}}
 const AutoReverseEnzyme = AutoEnzyme{<:ReverseMode}
 const AutoReverseOrNothingEnzyme = Union{AutoEnzyme{<:ReverseMode},AutoEnzyme{Nothing}}
 
-forward_mode(backend::AutoEnzyme{<:ForwardMode}) = backend.mode
+# forward mode if possible
+forward_mode(backend::AutoEnzyme{<:Mode}) = backend.mode
 forward_mode(::AutoEnzyme{Nothing}) = Forward
 
-reverse_mode(backend::AutoEnzyme{<:ReverseMode}) = backend.mode
+# reverse mode if possible
+reverse_mode(backend::AutoEnzyme{<:Mode}) = backend.mode
 reverse_mode(::AutoEnzyme{Nothing}) = Reverse
 
 DI.check_available(::AutoEnzyme) = true
@@ -60,10 +71,14 @@ function zero_sametype!(x_target, x)
     return x_sametype
 end
 
+include("utils.jl")
+
 include("forward_onearg.jl")
 include("forward_twoarg.jl")
 
 include("reverse_onearg.jl")
 include("reverse_twoarg.jl")
 
+include("common_onearg.jl")
+
 end # module

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/common_onearg.jl

@@ -0,0 +1,40 @@
+## Second derivative
+
+DI.prepare_second_derivative(f, ::AutoEnzyme, x, v) = NoSecondDerivativeExtras()
+
+function DI.second_derivative(
+    f, backend::AutoForwardOrNothingEnzyme, x, ::NoSecondDerivativeExtras
+)
+    df = DeferredDerivative(f, forward_mode(backend))
+    return DI.derivative(df, AutoEnzyme(forward_mode(backend)), x)
+end
+
+function DI.second_derivative!(
+    f, der2, backend::AutoForwardOrNothingEnzyme, x, ::NoSecondDerivativeExtras
+)
+    df = DeferredDerivative(f, forward_mode(backend))
+    return DI.derivative!(df, der2, AutoEnzyme(forward_mode(backend)), x)
+end
+
+## Hessian
+
+struct EnzymeHessianExtras{G,JE} <: HessianExtras
+    ∇f::G
+    jac_extras::JE
+end
+
+function DI.prepare_hessian(f, backend::AutoEnzyme, x)
+    ∇f = DeferredGradient(f, reverse_mode(backend))
+    jac_extras = DI.prepare_jacobian(∇f, AutoEnzyme(forward_mode(backend)), x)
+    return EnzymeHessianExtras(∇f, jac_extras)
+end
+
+function DI.hessian(f, backend::AutoEnzyme, x, extras::EnzymeHessianExtras)
+    @compat (; ∇f, jac_extras) = extras
+    return DI.jacobian(∇f, AutoEnzyme(forward_mode(backend)), x, jac_extras)
+end
+
+function DI.hessian!(f, hess, backend::AutoEnzyme, x, extras::EnzymeHessianExtras)
+    @compat (; ∇f, jac_extras) = extras
+    return DI.jacobian!(∇f, hess, AutoEnzyme(forward_mode(backend)), x, jac_extras)
+end

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl

@@ -6,7 +6,9 @@ function DI.value_and_pushforward(
     f, backend::AutoForwardOrNothingEnzyme, x, dx, ::NoPushforwardExtras
 )
     dx_sametype = convert(typeof(x), dx)
-    y, new_dy = autodiff(forward_mode(backend), f, Duplicated, Duplicated(x, dx_sametype))
+    y, new_dy = autodiff(
+        forward_mode(backend), Const(f), Duplicated, Duplicated(x, dx_sametype)
+    )
     return y, new_dy
 end
 
@@ -15,7 +17,9 @@ function DI.pushforward(
 )
     dx_sametype = convert(typeof(x), dx)
     new_dy = only(
-        autodiff(forward_mode(backend), f, DuplicatedNoNeed, Duplicated(x, dx_sametype))
+        autodiff(
+            forward_mode(backend), Const(f), DuplicatedNoNeed, Duplicated(x, dx_sametype)
+        ),
     )
     return new_dy
 end
@@ -121,3 +125,17 @@ function DI.value_and_jacobian!(
     y, new_jac = DI.value_and_jacobian(f, backend, x, extras)
     return y, copyto!(jac, new_jac)
 end
+
+## HVP
+
+DI.prepare_hvp(f, ::AutoForwardOrNothingEnzyme, x, v) = NoHVPExtras()
+
+function DI.hvp(f, backend::AutoForwardOrNothingEnzyme, x, v, ::NoHVPExtras)
+    ∇f = DeferredGradient(f, reverse_mode(backend))
+    return DI.pushforward(∇f, AutoEnzyme(forward_mode(backend)), x, v)
+end
+
+function DI.hvp!(f, p, backend::AutoForwardOrNothingEnzyme, x, v, ::NoHVPExtras)
+    ∇f = DeferredGradient(f, reverse_mode(backend))
+    return DI.pushforward!(∇f, p, AutoEnzyme(forward_mode(backend)), x, v)
+end

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_twoarg.jl

@@ -9,7 +9,7 @@ function DI.value_and_pushforward(
     dy_sametype = zero(y)
     autodiff(
         forward_mode(backend),
-        f!,
+        Const(f!),
         Const,
         Duplicated(y, dy_sametype),
         Duplicated(x, dx_sametype),

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl

@@ -7,7 +7,7 @@ DI.prepare_pullback(f, ::AutoReverseOrNothingEnzyme, x, dy) = NoPullbackExtras()
 function DI.value_and_pullback(
     f, ::AutoReverseOrNothingEnzyme, x::Number, dy::Number, ::NoPullbackExtras
 )
-    der, y = autodiff(ReverseWithPrimal, f, Active, Active(x))
+    der, y = autodiff(ReverseWithPrimal, Const(f), Active, Active(x))
     new_dx = dy * only(der)
     return y, new_dx
 end
@@ -43,7 +43,7 @@ function DI.value_and_pullback!(
     f, dx, ::AutoReverseOrNothingEnzyme, x::AbstractArray, dy::Number, ::NoPullbackExtras
 )
     dx_sametype = zero_sametype!(dx, x)
-    _, y = autodiff(ReverseWithPrimal, f, Active, Duplicated(x, dx_sametype))
+    _, y = autodiff(ReverseWithPrimal, Const(f), Active, Duplicated(x, dx_sametype))
     dx_sametype .*= dy
     return y, copyto!(dx, dx_sametype)
 end
@@ -155,3 +155,17 @@ function DI.value_and_jacobian!(
 end
 
 =#
+
+## HVP
+
+DI.prepare_hvp(f, ::AutoReverseEnzyme, x, v) = NoHVPExtras()
+
+function DI.hvp(f, backend::AutoReverseEnzyme, x, v, ::NoHVPExtras)
+    ∇f = DeferredGradient(f, reverse_mode(backend))
+    return DI.pullback(∇f, backend, x, v)
+end
+
+function DI.hvp!(f, p, backend::AutoReverseEnzyme, x, v, ::NoHVPExtras)
+    ∇f = DeferredGradient(f, reverse_mode(backend))
+    return DI.pullback!(∇f, p, backend, x, v)
+end

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl

@@ -0,0 +1,29 @@
+struct DeferredGradient{F,M<:Mode}
+    f::F
+    mode::M
+end
+
+function (def_grad::DeferredGradient{F,<:ForwardMode})(z) where {F}
+    return error("Not implemented yet")
+end
+
+function (def_grad::DeferredGradient{F,<:ReverseMode})(z) where {F}
+    @compat (; f, mode) = def_grad
+    grad = make_zero(z)
+    autodiff_deferred(mode, Const(f), Active, Duplicated(z, grad))
+    return grad
+end
+
+struct DeferredDerivative{F,M<:Mode}
+    f::F
+    mode::M
+end
+
+function (def_der::DeferredDerivative{F,<:ForwardMode})(z) where {F}
+    @compat (; f, mode) = def_der
+    return only(autodiff_deferred(mode, Const(f), DuplicatedNoNeed, Duplicated(z, one(z))))
+end
+
+function (def_der::DeferredDerivative{F,<:ReverseMode})(z) where {F}
+    return error("Not implemented yet")
+end

DifferentiationInterface/src/utils/check.jl

@@ -16,7 +16,7 @@ Check whether `backend` supports differentiation of two-argument functions.
 """
 check_twoarg(backend::AbstractADType) = Bool(twoarg_support(backend))
 
-sqnorm(x::AbstractArray) = sum(abs2, x)
+sqnorm(x::AbstractArray) = sum(abs2.(x))
 
 """
     check_hessian(backend)

DifferentiationInterface/src/utils/printing.jl

@@ -24,7 +24,7 @@ function backend_str(backend::AbstractADType)
     elseif mode(backend) isa SymbolicMode
         return "$bs (symbolic)"
     elseif mode(backend) isa ForwardOrReverseMode
-        return "$bs (forward/reverse)"
+        return "$bs (forward|reverse)"
     else
         error("Unknown mode")
     end

DifferentiationInterface/test/Single/Enzyme.jl

@@ -10,29 +10,31 @@ dense_backends = [
     AutoEnzyme(; mode=Enzyme.Reverse),
 ]
 
-sparse_backends = [
-    AutoSparse(
-        AutoEnzyme(; mode=Enzyme.Forward);
+sparse_backends =
+    AutoSparse.(
+        dense_backends,
         sparsity_detector=TracerSparsityDetector(),
         coloring_algorithm=GreedyColoringAlgorithm(),
-    ),
-    AutoSparse(
-        AutoEnzyme(; mode=Enzyme.Reverse);
-        sparsity_detector=TracerSparsityDetector(),
-        coloring_algorithm=GreedyColoringAlgorithm(),
-    ),
-]
+    )
 
 for backend in vcat(dense_backends, sparse_backends)
     @test check_available(backend)
     @test check_twoarg(backend)
-    @test !check_hessian(backend; verbose=false)
+    # @test !check_hessian(backend; verbose=false)
 end
 
 ## Dense backends
 
 test_differentiation(dense_backends; second_order=false, logging=LOGGING);
 
+test_differentiation(AutoEnzyme(); first_order=false, logging=LOGGING);
+test_differentiation(
+    AutoEnzyme(; mode=Enzyme.Forward);
+    first_order=false,
+    excluded=[HVPScenario, HessianScenario],
+    logging=LOGGING,
+);
+
 test_differentiation(
     AutoEnzyme(; mode=Enzyme.Forward);  # TODO: add more
     correctness=false,

DifferentiationInterfaceTest/src/scenarios/default.jl

@@ -1,6 +1,5 @@
 #=
 Constraints on the scenarios:
-- non-allocating whenever possible
 - type-stable
 - GPU-compatible (no scalar indexing)
 - vary shapes to be tricky