Introduce type hierarchy for extras (#127)

* Introduce type hierarchy for extras

* Fix docs

Author: gdalle

docs/src/core.md

@@ -14,41 +14,47 @@ DifferentiationInterface
 ```@autodocs
 Modules = [DifferentiationInterface]
 Pages = ["src/derivative.jl"]
+Private = false
 ```
 
 ## Gradient
 
 ```@autodocs
 Modules = [DifferentiationInterface]
 Pages = ["gradient.jl"]
+Private = false
 ```
 
 ## Jacobian
 
 ```@autodocs
 Modules = [DifferentiationInterface]
 Pages = ["jacobian.jl"]
+Private = false
 ```
 
 ## Second order
 
 ```@autodocs
 Modules = [DifferentiationInterface]
 Pages = ["second_order.jl", "second_derivative.jl", "hessian.jl", "hvp.jl"]
+Private = false
 ```
 
 ## Primitives
 
 ```@autodocs
 Modules = [DifferentiationInterface]
 Pages = ["pushforward.jl", "pullback.jl"]
+Private = false
 ```
 
 ## Backend queries
 
 ```@autodocs
 Modules = [DifferentiationInterface]
 Pages = ["backends.jl"]
+Private = false
 ```
 
 ## Internals
@@ -58,6 +64,4 @@ This is not part of the public API.
 ```@autodocs
 Modules = [DifferentiationInterface]
 Public = false
-Order = [:function, :type]
-Filter = t -> !(t isa Type && t <: ADTypes.AbstractADType)
 ```

ext/DifferentiationInterfaceChainRulesCoreExt/DifferentiationInterfaceChainRulesCoreExt.jl

@@ -4,6 +4,7 @@ using ADTypes: ADTypes, AutoChainRules
 using ChainRulesCore:
     HasForwardsMode, HasReverseMode, NoTangent, RuleConfig, frule_via_ad, rrule_via_ad
 import DifferentiationInterface as DI
+using DifferentiationInterface: NoPullbackExtras, NoPushforwardExtras
 
 ruleconfig(backend::AutoChainRules) = backend.ruleconfig
 
@@ -14,15 +15,23 @@ DI.supports_mutation(::AutoChainRules) = DI.MutationNotSupported()
 DI.mode(::AutoForwardChainRules) = ADTypes.AbstractForwardMode
 DI.mode(::AutoReverseChainRules) = ADTypes.AbstractReverseMode
 
-## Primitives
+## Pushforward
 
-function DI.value_and_pushforward(f, backend::AutoForwardChainRules, x, dx, extras::Nothing)
+DI.prepare_pushforward(f, ::AutoForwardChainRules, x) = NoPushforwardExtras()
+
+function DI.value_and_pushforward(
+    f, backend::AutoForwardChainRules, x, dx, ::NoPushforwardExtras
+)
     rc = ruleconfig(backend)
     y, new_dy = frule_via_ad(rc, (NoTangent(), dx), f, x)
     return y, new_dy
 end
 
-function DI.value_and_pullback(f, backend::AutoReverseChainRules, x, dy, extras::Nothing)
+## Pullback
+
+DI.prepare_pullback(f, ::AutoForwardChainRules, x) = NoPullbackExtras()
+
+function DI.value_and_pullback(f, backend::AutoReverseChainRules, x, dy, ::NoPullbackExtras)
     rc = ruleconfig(backend)
     y, pullback = rrule_via_ad(rc, f, x)
     _, new_dx = pullback(dy)

ext/DifferentiationInterfaceDiffractorExt/DifferentiationInterfaceDiffractorExt.jl

@@ -3,13 +3,18 @@ module DifferentiationInterfaceDiffractorExt
 import AbstractDifferentiation as AD  # public API for Diffractor
 using ADTypes: ADTypes, AutoChainRules, AutoDiffractor
 import DifferentiationInterface as DI
+using DifferentiationInterface: NoPushforwardExtras
 using Diffractor: DiffractorForwardBackend, DiffractorRuleConfig
 
 DI.supports_mutation(::AutoDiffractor) = DI.MutationNotSupported()
 DI.mode(::AutoDiffractor) = ADTypes.AbstractForwardMode
 DI.mode(::AutoChainRules{<:DiffractorRuleConfig}) = ADTypes.AbstractForwardMode
 
-function DI.value_and_pushforward(f, ::AutoDiffractor, x, dx, extras::Nothing)
+## Pushforward
+
+DI.prepare_pushforward(f, ::AutoDiffractor, x) = NoPushforwardExtras()
+
+function DI.value_and_pushforward(f, ::AutoDiffractor, x, dx, ::NoPushforwardExtras)
     vpff = AD.value_and_pushforward_function(DiffractorForwardBackend(), f, x)
     y, dy = vpff((dx,))
     return y, dy

ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl

@@ -2,6 +2,12 @@ module DifferentiationInterfaceEnzymeExt
 
 using ADTypes: ADTypes, AutoEnzyme
 import DifferentiationInterface as DI
+using DifferentiationInterface:
+    NoDerivativeExtras,
+    NoGradientExtras,
+    NoJacobianExtras,
+    NoPullbackExtras,
+    NoPushforwardExtras
 using DocStringExtensions
 using Enzyme:
     Active,

ext/DifferentiationInterfaceEnzymeExt/forward_allocating.jl

@@ -1,76 +1,86 @@
 ## Pushforward
 
-function DI.value_and_pushforward(f, backend::AutoForwardEnzyme, x, dx, extras::Nothing)
+DI.prepare_pushforward(f, ::AutoForwardEnzyme, x) = NoPushforwardExtras()
+
+function DI.value_and_pushforward(
+    f, backend::AutoForwardEnzyme, x, dx, ::NoPushforwardExtras
+)
     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)
+function DI.pushforward(f, backend::AutoForwardEnzyme, x, dx, ::NoPushforwardExtras)
     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
+    f, _dy, backend::AutoForwardEnzyme, x, dx, extras::NoPushforwardExtras
 )
     # 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)
+function DI.pushforward!!(
+    f, _dy, backend::AutoForwardEnzyme, x, dx, extras::NoPushforwardExtras
+)
     # 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)
+DI.prepare_gradient(f, ::AutoForwardEnzyme, x) = NoGradientExtras()
+
+function DI.gradient(f, backend::AutoForwardEnzyme, x::AbstractArray, ::NoGradientExtras)
     return reshape(collect(gradient(backend.mode, f, x)), size(x))
 end
 
 function DI.value_and_gradient(
-    f, backend::AutoForwardEnzyme, x::AbstractArray, extras::Nothing
+    f, backend::AutoForwardEnzyme, x::AbstractArray, extras::NoGradientExtras
 )
     return f(x), DI.gradient(f, backend, x, extras)
 end
 
 function DI.gradient!!(
-    f, _grad, backend::AutoForwardEnzyme, x::AbstractArray, extras::Nothing
+    f, _grad, backend::AutoForwardEnzyme, x::AbstractArray, extras::NoGradientExtras
 )
     return DI.gradient(f, backend, x, extras)
 end
 
 function DI.value_and_gradient!!(
-    f, _grad, backend::AutoForwardEnzyme, x::AbstractArray, extras::Nothing
+    f, _grad, backend::AutoForwardEnzyme, x::AbstractArray, extras::NoGradientExtras
 )
     return DI.value_and_gradient(f, backend, x, extras)
 end
 
 ## Jacobian
 
-function DI.jacobian(f, backend::AutoForwardEnzyme, x::AbstractArray, extras::Nothing)
+DI.prepare_jacobian(f, ::AutoForwardEnzyme, x) = NoJacobianExtras()
+
+function DI.jacobian(f, backend::AutoForwardEnzyme, x::AbstractArray, ::NoJacobianExtras)
     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
+    f, backend::AutoForwardEnzyme, x::AbstractArray, extras::NoJacobianExtras
 )
     return f(x), DI.jacobian(f, backend, x, extras)
 end
 
 function DI.jacobian!!(
-    f, _jac, backend::AutoForwardEnzyme, x::AbstractArray, extras::Nothing
+    f, _jac, backend::AutoForwardEnzyme, x::AbstractArray, extras::NoJacobianExtras
 )
     return DI.jacobian(f, backend, x, extras)
 end
 
 function DI.value_and_jacobian!!(
-    f, _jac, backend::AutoForwardEnzyme, x::AbstractArray, extras::Nothing
+    f, _jac, backend::AutoForwardEnzyme, x::AbstractArray, extras::NoJacobianExtras
 )
     return DI.value_and_jacobian(f, backend, x, extras)
 end

ext/DifferentiationInterfaceEnzymeExt/forward_mutating.jl

@@ -1,7 +1,9 @@
 ## Pushforward
 
+DI.prepare_pushforward(f!, ::AutoForwardEnzyme, y, x) = NoPushforwardExtras()
+
 function DI.value_and_pushforward!!(
-    f!, y, dy, backend::AutoForwardEnzyme, x, dx, extras::Nothing
+    f!, y, dy, backend::AutoForwardEnzyme, x, dx, ::NoPushforwardExtras
 )
     dx_sametype = convert(typeof(x), dx)
     dy_sametype = zero_sametype!!(dy, y)

ext/DifferentiationInterfaceEnzymeExt/reverse_allocating.jl

@@ -1,17 +1,19 @@
 ## Pullback
 
+DI.prepare_pullback(f, ::AutoReverseEnzyme, x) = NoPullbackExtras()
+
 ### Out-of-place
 
 function DI.value_and_pullback(
-    f, ::AutoReverseEnzyme, x::Number, dy::Number, extras::Nothing
+    f, ::AutoReverseEnzyme, x::Number, dy::Number, ::NoPullbackExtras
 )
     der, y = autodiff(ReverseWithPrimal, f, Active, Active(x))
     new_dx = dy * only(der)
     return y, new_dx
 end
 
 function DI.value_and_pullback(
-    f, ::AutoReverseEnzyme, x::Number, dy::AbstractArray, extras::Nothing
+    f, ::AutoReverseEnzyme, x::Number, dy::AbstractArray, ::NoPullbackExtras
 )
     forw, rev = autodiff_thunk(
         ReverseSplitWithPrimal, Const{typeof(f)}, Duplicated, Active{typeof(x)}
@@ -25,7 +27,7 @@ end
 ### In-place
 
 function DI.value_and_pullback!!(
-    f, dx, ::AutoReverseEnzyme, x::AbstractArray, dy::Number, extras::Nothing
+    f, dx, ::AutoReverseEnzyme, x::AbstractArray, dy::Number, ::NoPullbackExtras
 )
     dx_sametype = zero_sametype!!(dx, x)
     _, y = autodiff(ReverseWithPrimal, f, Active, Duplicated(x, dx_sametype))
@@ -34,7 +36,7 @@ function DI.value_and_pullback!!(
 end
 
 function DI.value_and_pullback!!(
-    f, dx, ::AutoReverseEnzyme, x::AbstractArray, dy::AbstractArray, extras::Nothing
+    f, dx, ::AutoReverseEnzyme, x::AbstractArray, dy::AbstractArray, ::NoPullbackExtras
 )
     dx_sametype = zero_sametype!!(dx, x)
     forw, rev = autodiff_thunk(
@@ -46,18 +48,22 @@ function DI.value_and_pullback!!(
     return y, dx_sametype
 end
 
-function DI.value_and_pullback(f, backend::AutoReverseEnzyme, x::AbstractArray, dy, extras)
+function DI.value_and_pullback(
+    f, backend::AutoReverseEnzyme, x::AbstractArray, dy, extras::NoPullbackExtras
+)
     dx = similar(x)
     return DI.value_and_pullback!!(f, dx, backend, x, dy, extras)
 end
 
 ## Gradient
 
-function DI.gradient(f, ::AutoReverseEnzyme, x::AbstractArray, extras::Nothing)
+DI.prepare_gradient(f, ::AutoReverseEnzyme) = NoGradientExtras()
+
+function DI.gradient(f, ::AutoReverseEnzyme, x::AbstractArray, ::NoGradientExtras)
     return gradient(Reverse, f, x)
 end
 
-function DI.gradient!!(f, grad, ::AutoReverseEnzyme, x::AbstractArray, extras::Nothing)
+function DI.gradient!!(f, grad, ::AutoReverseEnzyme, x::AbstractArray, ::NoGradientExtras)
     grad_sametype = convert(typeof(x), grad)
     gradient!(Reverse, grad_sametype, f, x)
     return grad_sametype

ext/DifferentiationInterfaceEnzymeExt/reverse_mutating.jl

@@ -1,15 +1,17 @@
 ## Pullback
 
+DI.prepare_pullback(f!, ::AutoReverseEnzyme, y, x) = NoPullbackExtras()
+
 function DI.value_and_pullback!!(
-    f!, y, _dx, ::AutoReverseEnzyme, x::Number, dy, extras::Nothing
+    f!, y, _dx, ::AutoReverseEnzyme, x::Number, dy, ::NoPullbackExtras
 )
     dy_sametype = convert(typeof(y), copy(dy))
     _, new_dx = only(autodiff(Reverse, f!, Const, Duplicated(y, dy_sametype), Active(x)))
     return y, new_dx
 end
 
 function DI.value_and_pullback!!(
-    f!, y, dx, ::AutoReverseEnzyme, x::AbstractArray, dy, extras::Nothing
+    f!, y, dx, ::AutoReverseEnzyme, x::AbstractArray, dy, ::NoPullbackExtras
 )
     dx_sametype = zero_sametype!!(dx, x)
     dy_sametype = convert(typeof(y), copy(dy))

ext/DifferentiationInterfaceFastDifferentiationExt/DifferentiationInterfaceFastDifferentiationExt.jl

@@ -1,8 +1,16 @@
 module DifferentiationInterfaceFastDifferentiationExt
 
 using ADTypes: ADTypes
-using DifferentiationInterface: AutoFastDifferentiation, AutoSparseFastDifferentiation
 import DifferentiationInterface as DI
+using DifferentiationInterface: AutoFastDifferentiation, AutoSparseFastDifferentiation
+using DifferentiationInterface:
+    DerivativeExtras,
+    GradientExtras,
+    HessianExtras,
+    HVPExtras,
+    JacobianExtras,
+    PullbackExtras,
+    PushforwardExtras
 using FastDifferentiation:
     derivative,
     hessian,