Only jacobian and hessian for sparse backends (#179)

* Only jacobian and hessian for sparse backends

* Put sparse scenarios with the other scenarios

Author: gdalle

DifferentiationInterface/docs/src/backends.md

@@ -63,6 +63,9 @@ AutoZygote
 
 ### Sparse
 
+!!! warning
+    For sparse backends, only the Jacobian and Hessian operators are implemented.
+
 ```@docs
 AutoSparseFastDifferentiation
 AutoSparseFiniteDiff

DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/DifferentiationInterfaceChainRulesCoreExt.jl

@@ -11,6 +11,7 @@ ruleconfig(backend::AutoChainRules) = backend.ruleconfig
 const AutoForwardChainRules = AutoChainRules{<:RuleConfig{>:HasForwardsMode}}
 const AutoReverseChainRules = AutoChainRules{<:RuleConfig{>:HasReverseMode}}
 
+DI.check_available(::AutoChainRules) = true
 DI.supports_mutation(::AutoChainRules) = DI.MutationNotSupported()
 DI.mode(::AutoForwardChainRules) = ADTypes.AbstractForwardMode
 DI.mode(::AutoReverseChainRules) = ADTypes.AbstractReverseMode

DifferentiationInterface/ext/DifferentiationInterfaceDiffractorExt/DifferentiationInterfaceDiffractorExt.jl

@@ -5,6 +5,7 @@ import DifferentiationInterface as DI
 using DifferentiationInterface: NoPushforwardExtras
 using Diffractor: DiffractorRuleConfig, TaylorTangentIndex, ZeroBundle, bundle, ∂☆
 
+DI.check_available(::AutoDiffractor) = true
 DI.supports_mutation(::AutoDiffractor) = DI.MutationNotSupported()
 DI.mode(::AutoDiffractor) = ADTypes.AbstractForwardMode
 DI.mode(::AutoChainRules{<:DiffractorRuleConfig}) = ADTypes.AbstractForwardMode

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl

@@ -38,6 +38,8 @@ AutoEnzyme
 const AutoForwardEnzyme = AutoEnzyme{<:ForwardMode}
 const AutoReverseEnzyme = AutoEnzyme{<:ReverseMode}
 
+DI.check_available(::AutoEnzyme) = true
+
 function DI.mode(::AutoEnzyme)
     return error(
         "You need to specify the Enzyme mode with `AutoEnzyme(Enzyme.Forward)` or `AutoEnzyme(Enzyme.Reverse)`",

DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/DifferentiationInterfaceFastDifferentiationExt.jl

@@ -31,6 +31,7 @@ const AnyAutoFastDifferentiation = Union{
     AutoFastDifferentiation,AutoSparseFastDifferentiation
 }
 
+DI.check_available(::AutoFastDifferentiation) = true
 DI.mode(::AnyAutoFastDifferentiation) = ADTypes.AbstractSymbolicDifferentiationMode
 DI.pushforward_performance(::AnyAutoFastDifferentiation) = DI.PushforwardFast()
 DI.pullback_performance(::AnyAutoFastDifferentiation) = DI.PullbackSlow()

DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/DifferentiationInterfaceFiniteDiffExt.jl

@@ -23,7 +23,7 @@ using FiniteDiff:
     finite_difference_jacobian!
 using LinearAlgebra: dot, mul!
 
-const AnyAutoFiniteDiff = Union{AutoFiniteDiff,AutoSparseFiniteDiff}
+DI.check_available(::AutoFiniteDiff) = true
 
 # see https://github.com/SciML/ADTypes.jl/issues/33
 

DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/onearg.jl

@@ -1,16 +1,14 @@
 ## Pushforward
 
-DI.prepare_pushforward(f, ::AnyAutoFiniteDiff, x) = NoPushforwardExtras()
+DI.prepare_pushforward(f, ::AutoFiniteDiff, x) = NoPushforwardExtras()
 
-function DI.pushforward(f, backend::AnyAutoFiniteDiff, x, dx, ::NoPushforwardExtras)
+function DI.pushforward(f, backend::AutoFiniteDiff, x, dx, ::NoPushforwardExtras)
     step(t::Number) = f(x .+ t .* dx)
     new_dy = finite_difference_derivative(step, zero(eltype(x)), fdtype(backend))
     return new_dy
 end
 
-function DI.value_and_pushforward(
-    f, backend::AnyAutoFiniteDiff, x, dx, ::NoPushforwardExtras
-)
+function DI.value_and_pushforward(f, backend::AutoFiniteDiff, x, dx, ::NoPushforwardExtras)
     y = f(x)
     step(t::Number) = f(x .+ t .* dx)
     new_dy = finite_difference_derivative(
@@ -25,7 +23,7 @@ struct FiniteDiffOneArgDerivativeExtras{C}
     cache::C
 end
 
-function DI.prepare_derivative(f, backend::AnyAutoFiniteDiff, x)
+function DI.prepare_derivative(f, backend::AutoFiniteDiff, x)
     y = f(x)
     cache = if y isa Number
         nothing
@@ -39,26 +37,26 @@ end
 ### Scalar to scalar
 
 function DI.derivative(
-    f, backend::AnyAutoFiniteDiff, x, ::FiniteDiffOneArgDerivativeExtras{Nothing}
+    f, backend::AutoFiniteDiff, x, ::FiniteDiffOneArgDerivativeExtras{Nothing}
 )
     return finite_difference_derivative(f, x, fdtype(backend))
 end
 
 function DI.derivative!(
-    f, der, backend::AnyAutoFiniteDiff, x, extras::FiniteDiffOneArgDerivativeExtras{Nothing}
+    f, der, backend::AutoFiniteDiff, x, extras::FiniteDiffOneArgDerivativeExtras{Nothing}
 )
     return copyto!(der, DI.derivative(f, backend, x, extras))
 end
 
 function DI.value_and_derivative(
-    f, backend::AnyAutoFiniteDiff, x, ::FiniteDiffOneArgDerivativeExtras{Nothing}
+    f, backend::AutoFiniteDiff, x, ::FiniteDiffOneArgDerivativeExtras{Nothing}
 )
     y = f(x)
     return y, finite_difference_derivative(f, x, fdtype(backend), eltype(y), y)
 end
 
 function DI.value_and_derivative!(
-    f, der, backend::AnyAutoFiniteDiff, x, extras::FiniteDiffOneArgDerivativeExtras{Nothing}
+    f, der, backend::AutoFiniteDiff, x, extras::FiniteDiffOneArgDerivativeExtras{Nothing}
 )
     y, new_der = DI.value_and_derivative(f, backend, x, extras)
     return y, copyto!(der, new_der)
@@ -67,34 +65,26 @@ end
 ### Scalar to array
 
 function DI.derivative(
-    f, ::AnyAutoFiniteDiff, x, extras::FiniteDiffOneArgDerivativeExtras{<:GradientCache}
+    f, ::AutoFiniteDiff, x, extras::FiniteDiffOneArgDerivativeExtras{<:GradientCache}
 )
     return finite_difference_gradient(f, x, extras.cache)
 end
 
 function DI.derivative!(
-    f,
-    der,
-    ::AnyAutoFiniteDiff,
-    x,
-    extras::FiniteDiffOneArgDerivativeExtras{<:GradientCache},
+    f, der, ::AutoFiniteDiff, x, extras::FiniteDiffOneArgDerivativeExtras{<:GradientCache}
 )
     return finite_difference_gradient!(der, f, x, extras.cache)
 end
 
 function DI.value_and_derivative(
-    f, ::AnyAutoFiniteDiff, x, extras::FiniteDiffOneArgDerivativeExtras{<:GradientCache}
+    f, ::AutoFiniteDiff, x, extras::FiniteDiffOneArgDerivativeExtras{<:GradientCache}
 )
     y = f(x)
     return y, finite_difference_gradient(f, x, extras.cache)
 end
 
 function DI.value_and_derivative!(
-    f,
-    der,
-    ::AnyAutoFiniteDiff,
-    x,
-    extras::FiniteDiffOneArgDerivativeExtras{<:GradientCache},
+    f, der, ::AutoFiniteDiff, x, extras::FiniteDiffOneArgDerivativeExtras{<:GradientCache}
 )
     return f(x), finite_difference_gradient!(der, f, x, extras.cache)
 end
@@ -105,33 +95,33 @@ struct FiniteDiffGradientExtras{C}
     cache::C
 end
 
-function DI.prepare_gradient(f, backend::AnyAutoFiniteDiff, x)
+function DI.prepare_gradient(f, backend::AutoFiniteDiff, x)
     y = f(x)
     df = zero(y) .* x
     cache = GradientCache(df, x, fdtype(backend))
     return FiniteDiffGradientExtras(cache)
 end
 
 function DI.gradient(
-    f, ::AnyAutoFiniteDiff, x::AbstractArray, extras::FiniteDiffGradientExtras
+    f, ::AutoFiniteDiff, x::AbstractArray, extras::FiniteDiffGradientExtras
 )
     return finite_difference_gradient(f, x, extras.cache)
 end
 
 function DI.value_and_gradient(
-    f, ::AnyAutoFiniteDiff, x::AbstractArray, extras::FiniteDiffGradientExtras
+    f, ::AutoFiniteDiff, x::AbstractArray, extras::FiniteDiffGradientExtras
 )
     return f(x), finite_difference_gradient(f, x, extras.cache)
 end
 
 function DI.gradient!(
-    f, grad, ::AnyAutoFiniteDiff, x::AbstractArray, extras::FiniteDiffGradientExtras
+    f, grad, ::AutoFiniteDiff, x::AbstractArray, extras::FiniteDiffGradientExtras
 )
     return finite_difference_gradient!(grad, f, x, extras.cache)
 end
 
 function DI.value_and_gradient!(
-    f, grad, ::AnyAutoFiniteDiff, x::AbstractArray, extras::FiniteDiffGradientExtras
+    f, grad, ::AutoFiniteDiff, x::AbstractArray, extras::FiniteDiffGradientExtras
 )
     return f(x), finite_difference_gradient!(grad, f, x, extras.cache)
 end
@@ -142,7 +132,7 @@ struct FiniteDiffOneArgJacobianExtras{C}
     cache::C
 end
 
-function DI.prepare_jacobian(f, backend::AnyAutoFiniteDiff, x)
+function DI.prepare_jacobian(f, backend::AutoFiniteDiff, x)
     y = f(x)
     x1 = similar(x)
     fx = similar(y)
@@ -151,25 +141,23 @@ function DI.prepare_jacobian(f, backend::AnyAutoFiniteDiff, x)
     return FiniteDiffOneArgJacobianExtras(cache)
 end
 
-function DI.jacobian(f, ::AnyAutoFiniteDiff, x, extras::FiniteDiffOneArgJacobianExtras)
+function DI.jacobian(f, ::AutoFiniteDiff, x, extras::FiniteDiffOneArgJacobianExtras)
     return finite_difference_jacobian(f, x, extras.cache)
 end
 
 function DI.value_and_jacobian(
-    f, ::AnyAutoFiniteDiff, x, extras::FiniteDiffOneArgJacobianExtras
+    f, ::AutoFiniteDiff, x, extras::FiniteDiffOneArgJacobianExtras
 )
     y = f(x)
     return y, finite_difference_jacobian(f, x, extras.cache, y)
 end
 
-function DI.jacobian!(
-    f, jac, ::AnyAutoFiniteDiff, x, extras::FiniteDiffOneArgJacobianExtras
-)
+function DI.jacobian!(f, jac, ::AutoFiniteDiff, x, extras::FiniteDiffOneArgJacobianExtras)
     return copyto!(jac, finite_difference_jacobian(f, x, extras.cache; jac_prototype=jac))
 end
 
 function DI.value_and_jacobian!(
-    f, jac, ::AnyAutoFiniteDiff, x, extras::FiniteDiffOneArgJacobianExtras
+    f, jac, ::AutoFiniteDiff, x, extras::FiniteDiffOneArgJacobianExtras
 )
     y = f(x)
     return y,
@@ -182,15 +170,15 @@ struct FiniteDiffHessianExtras{C}
     cache::C
 end
 
-function DI.prepare_hessian(f, backend::AnyAutoFiniteDiff, x)
+function DI.prepare_hessian(f, backend::AutoFiniteDiff, x)
     cache = HessianCache(x, fdhtype(backend))
     return FiniteDiffHessianExtras(cache)
 end
 
-function DI.hessian(f, ::AnyAutoFiniteDiff, x, extras::FiniteDiffHessianExtras)
+function DI.hessian(f, ::AutoFiniteDiff, x, extras::FiniteDiffHessianExtras)
     return finite_difference_hessian(f, x, extras.cache)
 end
 
-function DI.hessian!(f, hess, ::AnyAutoFiniteDiff, x, extras::FiniteDiffHessianExtras)
+function DI.hessian!(f, hess, ::AutoFiniteDiff, x, extras::FiniteDiffHessianExtras)
     return finite_difference_hessian!(hess, f, x, extras.cache)
 end

DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/twoarg.jl

@@ -1,9 +1,9 @@
 ## Pushforward
 
-DI.prepare_pushforward(f!, y, ::AnyAutoFiniteDiff, x) = NoPushforwardExtras()
+DI.prepare_pushforward(f!, y, ::AutoFiniteDiff, x) = NoPushforwardExtras()
 
 function DI.value_and_pushforward(
-    f!, y, backend::AnyAutoFiniteDiff, x, dx, ::NoPushforwardExtras
+    f!, y, backend::AutoFiniteDiff, x, dx, ::NoPushforwardExtras
 )
     function step(t::Number)::AbstractArray
         new_y = similar(y)
@@ -23,37 +23,37 @@ struct FiniteDiffTwoArgDerivativeExtras{C}
     cache::C
 end
 
-function DI.prepare_derivative(f!, y, ::AnyAutoFiniteDiff, x)
+function DI.prepare_derivative(f!, y, ::AutoFiniteDiff, x)
     cache = nothing
     return FiniteDiffTwoArgDerivativeExtras(cache)
 end
 
 function DI.value_and_derivative(
-    f!, y, backend::AnyAutoFiniteDiff, x, ::FiniteDiffTwoArgDerivativeExtras
+    f!, y, backend::AutoFiniteDiff, x, ::FiniteDiffTwoArgDerivativeExtras
 )
     f!(y, x)
     der = finite_difference_gradient(f!, x, fdtype(backend), eltype(y), FUNCTION_INPLACE, y)
     return y, der
 end
 
 function DI.value_and_derivative!(
-    f!, y, der, backend::AnyAutoFiniteDiff, x, ::FiniteDiffTwoArgDerivativeExtras
+    f!, y, der, backend::AutoFiniteDiff, x, ::FiniteDiffTwoArgDerivativeExtras
 )
     f!(y, x)
     finite_difference_gradient!(der, f!, x, fdtype(backend), eltype(y), FUNCTION_INPLACE, y)
     return y, der
 end
 
 function DI.derivative(
-    f!, y, backend::AnyAutoFiniteDiff, x, ::FiniteDiffTwoArgDerivativeExtras
+    f!, y, backend::AutoFiniteDiff, x, ::FiniteDiffTwoArgDerivativeExtras
 )
     f!(y, x)
     der = finite_difference_gradient(f!, x, fdtype(backend), eltype(y), FUNCTION_INPLACE, y)
     return der
 end
 
 function DI.derivative!(
-    f!, y, der, backend::AnyAutoFiniteDiff, x, ::FiniteDiffTwoArgDerivativeExtras
+    f!, y, der, backend::AutoFiniteDiff, x, ::FiniteDiffTwoArgDerivativeExtras
 )
     finite_difference_gradient!(der, f!, x, fdtype(backend), eltype(y), FUNCTION_INPLACE)
     return der
@@ -65,7 +65,7 @@ struct FiniteDiffTwoArgJacobianExtras{C}
     cache::C
 end
 
-function DI.prepare_jacobian(f!, y, backend::AnyAutoFiniteDiff, x)
+function DI.prepare_jacobian(f!, y, backend::AutoFiniteDiff, x)
     x1 = similar(x)
     fx = similar(y)
     fx1 = similar(y)
@@ -74,7 +74,7 @@ function DI.prepare_jacobian(f!, y, backend::AnyAutoFiniteDiff, x)
 end
 
 function DI.value_and_jacobian(
-    f!, y, ::AnyAutoFiniteDiff, x, extras::FiniteDiffTwoArgJacobianExtras
+    f!, y, ::AutoFiniteDiff, x, extras::FiniteDiffTwoArgJacobianExtras
 )
     jac = similar(y, length(y), length(x))
     finite_difference_jacobian!(jac, f!, x, extras.cache)
@@ -83,21 +83,21 @@ function DI.value_and_jacobian(
 end
 
 function DI.value_and_jacobian!(
-    f!, y, jac, ::AnyAutoFiniteDiff, x, extras::FiniteDiffTwoArgJacobianExtras
+    f!, y, jac, ::AutoFiniteDiff, x, extras::FiniteDiffTwoArgJacobianExtras
 )
     finite_difference_jacobian!(jac, f!, x, extras.cache)
     f!(y, x)
     return y, jac
 end
 
-function DI.jacobian(f!, y, ::AnyAutoFiniteDiff, x, extras::FiniteDiffTwoArgJacobianExtras)
+function DI.jacobian(f!, y, ::AutoFiniteDiff, x, extras::FiniteDiffTwoArgJacobianExtras)
     jac = similar(y, length(y), length(x))
     finite_difference_jacobian!(jac, f!, x, extras.cache)
     return jac
 end
 
 function DI.jacobian!(
-    f!, y, jac, ::AnyAutoFiniteDiff, x, extras::FiniteDiffTwoArgJacobianExtras
+    f!, y, jac, ::AutoFiniteDiff, x, extras::FiniteDiffTwoArgJacobianExtras
 )
     finite_difference_jacobian!(jac, f!, x, extras.cache)
     return jac

DifferentiationInterface/ext/DifferentiationInterfaceFiniteDifferencesExt/DifferentiationInterfaceFiniteDifferencesExt.jl

@@ -8,6 +8,7 @@ using FillArrays: OneElement
 using FiniteDifferences: FiniteDifferences, grad, jacobian, jvp, j′vp
 using LinearAlgebra: dot
 
+DI.check_available(::AutoFiniteDifferences) = true
 DI.supports_mutation(::AutoFiniteDifferences) = DI.MutationNotSupported()
 
 function FiniteDifferences.to_vec(a::OneElement)  # TODO: remove type piracy (https://github.com/JuliaDiff/FiniteDifferences.jl/issues/141)

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/DifferentiationInterfaceForwardDiffExt.jl

@@ -32,7 +32,7 @@ using ForwardDiff:
     value
 using LinearAlgebra: dot, mul!
 
-const AnyAutoForwardDiff{C,T} = Union{AutoForwardDiff{C,T},AutoSparseForwardDiff{C,T}}
+DI.check_available(::AutoForwardDiff) = true
 
 include("utils.jl")
 include("onearg.jl")