fix: clarify dispatch for preparation (#746)

* fix: separe `prepare` from the hidden `prepare_nokwarg`

* DOcs

* Typing

* Fix

* Toggle fail fast

Author: gdalle

DifferentiationInterface/docs/src/dev_guide.md

@@ -23,10 +23,10 @@ Most operators have 4 variants, which look like this in the first order: `operat
 To implement a new operator for an existing backend, you need to write 5 methods: 1 for [preparation](@ref Preparation) and 4 corresponding to the variants of the operator (see above).
 For first-order operators, you may also want to support [in-place functions](@ref "Mutation and signatures"), which requires another 5 methods (defined on `f!` instead of `f`).
 
-The method `prepare_operator` must output a `prep` object of the correct type.
-For instance, `prepare_gradient(f, backend, x)` must return a [`DifferentiationInterface.GradientPrep`](@ref).
-Assuming you don't need any preparation for said operator, you can use the trivial prep that are already defined, like `DifferentiationInterface.NoGradientPrep`.
-Otherwise, define a custom struct like `MyGradientPrep <: DifferentiationInterface.GradientPrep` and put the necessary storage in there.
+The method `prepare_operator_nokwarg` must output a `prep` object of the correct type.
+For instance, `prepare_gradient(strict, f, backend, x)` must return a [`DifferentiationInterface.GradientPrep`](@ref).
+Assuming you don't need any preparation for said operator, you can use the trivial prep that are already defined, like `DifferentiationInterface.NoGradientPrep{SIG}`.
+Otherwise, define a custom struct like `MyGradientPrep{SIG} <: DifferentiationInterface.GradientPrep{SIG}` and put the necessary storage in there.
 
 ## New backend
 
@@ -75,4 +75,4 @@ GROUP = get(ENV, "JULIA_DI_TEST_GROUP", "Back/SuperDiff")
 
 but don't forget to switch it back before pushing.
 
-Finally, you need to add your backend to the documentation, modifying every page that involves a list of backends (including the `README.md`).
+Finally, you need to add your backend to the documentation, modifying every page that involves a list of backends (including the `README.md`).

DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/differentiate_with.jl

@@ -1,7 +1,7 @@
 function ChainRulesCore.rrule(dw::DI.DifferentiateWith, x)
     (; f, backend) = dw
     y = f(x)
-    prep_same = DI.prepare_pullback_same_point(Val(true), f, backend, x, (y,))
+    prep_same = DI.prepare_pullback_same_point_nokwarg(Val(true), f, backend, x, (y,))
     function pullbackfunc(dy)
         tx = DI.pullback(f, prep_same, backend, x, (dy,))
         return (NoTangent(), only(tx))

DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/reverse_onearg.jl

@@ -6,7 +6,7 @@ struct ChainRulesPullbackPrepSamePoint{SIG,Y,PB} <: DI.PullbackPrep{SIG}
     pb::PB
 end
 
-function DI.prepare_pullback(
+function DI.prepare_pullback_nokwarg(
     strict::Val,
     f,
     backend::AutoReverseChainRules,

DifferentiationInterface/ext/DifferentiationInterfaceDiffractorExt/DifferentiationInterfaceDiffractorExt.jl

@@ -10,7 +10,9 @@ DI.pullback_performance(::AutoDiffractor) = DI.PullbackSlow()
 
 ## Pushforward
 
-function DI.prepare_pushforward(strict::Val, f, backend::AutoDiffractor, x, tx::NTuple)
+function DI.prepare_pushforward_nokwarg(
+    strict::Val, f, backend::AutoDiffractor, x, tx::NTuple
+)
     _sig = DI.signature(f, backend, x, tx; strict)
     return DI.NoPushforwardPrep(_sig)
 end

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl

@@ -1,6 +1,6 @@
 ## Pushforward
 
-function DI.prepare_pushforward(
+function DI.prepare_pushforward_nokwarg(
     strict::Val,
     f::F,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
@@ -122,7 +122,7 @@ struct EnzymeForwardGradientPrep{SIG,B,O} <: DI.GradientPrep{SIG}
     shadows::O
 end
 
-function DI.prepare_gradient(
+function DI.prepare_gradient_nokwarg(
     strict::Val,
     f::F,
     backend::AutoEnzyme{<:ForwardMode,<:Union{Nothing,Const}},
@@ -203,7 +203,7 @@ struct EnzymeForwardOneArgJacobianPrep{SIG,B,O} <: DI.JacobianPrep{SIG}
     output_length::Int
 end
 
-function DI.prepare_jacobian(
+function DI.prepare_jacobian_nokwarg(
     strict::Val,
     f::F,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing},<:Union{Nothing,Const}},

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_twoarg.jl

@@ -1,6 +1,6 @@
 ## Pushforward
 
-function DI.prepare_pushforward(
+function DI.prepare_pushforward_nokwarg(
     strict::Val,
     f!::F,
     y,

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl

@@ -52,7 +52,7 @@ struct EnzymeReverseOneArgPullbackPrep{SIG,Y} <: DI.PullbackPrep{SIG}
     y_example::Y  # useful to create return activity
 end
 
-function DI.prepare_pullback(
+function DI.prepare_pullback_nokwarg(
     strict::Val,
     f::F,
     backend::AutoEnzyme{<:Union{ReverseMode,Nothing}},
@@ -191,7 +191,7 @@ end
 
 ## Gradient
 
-function DI.prepare_gradient(
+function DI.prepare_gradient_nokwarg(
     strict::Val,
     f::F,
     backend::AutoEnzyme{<:Union{ReverseMode,Nothing}},

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_twoarg.jl

@@ -5,7 +5,7 @@ struct EnzymeReverseTwoArgPullbackPrep{SIG,TY} <: DI.PullbackPrep{SIG}
     ty_copy::TY
 end
 
-function DI.prepare_pullback(
+function DI.prepare_pullback_nokwarg(
     strict::Val,
     f!::F,
     y,

DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/onearg.jl

@@ -7,7 +7,7 @@ struct FastDifferentiationOneArgPushforwardPrep{SIG,Y,E1,E1!} <: DI.PushforwardP
     jvp_exe!::E1!
 end
 
-function DI.prepare_pushforward(
+function DI.prepare_pushforward_nokwarg(
     strict::Val,
     f,
     backend::AutoFastDifferentiation,
@@ -105,7 +105,7 @@ struct FastDifferentiationOneArgPullbackPrep{SIG,E1,E1!} <: DI.PullbackPrep{SIG}
     vjp_exe!::E1!
 end
 
-function DI.prepare_pullback(
+function DI.prepare_pullback_nokwarg(
     strict::Val,
     f,
     backend::AutoFastDifferentiation,
@@ -204,7 +204,7 @@ struct FastDifferentiationOneArgDerivativePrep{SIG,Y,E1,E1!} <: DI.DerivativePre
     der_exe!::E1!
 end
 
-function DI.prepare_derivative(
+function DI.prepare_derivative_nokwarg(
     strict::Val, f, backend::AutoFastDifferentiation, x, contexts::Vararg{DI.Context,C};
 ) where {C}
     _sig = DI.signature(f, backend, x, contexts...; strict)
@@ -284,7 +284,7 @@ struct FastDifferentiationOneArgGradientPrep{SIG,E1,E1!} <: DI.GradientPrep{SIG}
     jac_exe!::E1!
 end
 
-function DI.prepare_gradient(
+function DI.prepare_gradient_nokwarg(
     strict::Val, f, backend::AutoFastDifferentiation, x, contexts::Vararg{DI.Context,C};
 ) where {C}
     _sig = DI.signature(f, backend, x, contexts...; strict)
@@ -360,7 +360,7 @@ struct FastDifferentiationOneArgJacobianPrep{SIG,Y,E1,E1!} <: DI.JacobianPrep{SI
     jac_exe!::E1!
 end
 
-function DI.prepare_jacobian(
+function DI.prepare_jacobian_nokwarg(
     strict::Val,
     f,
     backend::Union{AutoFastDifferentiation,AutoSparse{<:AutoFastDifferentiation}},
@@ -445,7 +445,7 @@ struct FastDifferentiationAllocatingSecondDerivativePrep{SIG,Y,D,E2,E2!} <:
     der2_exe!::E2!
 end
 
-function DI.prepare_second_derivative(
+function DI.prepare_second_derivative_nokwarg(
     strict::Val, f, backend::AutoFastDifferentiation, x, contexts::Vararg{DI.Context,C};
 ) where {C}
     _sig = DI.signature(f, backend, x, contexts...; strict)
@@ -462,7 +462,7 @@ function DI.prepare_second_derivative(
     der2_exe = make_function(der2_vec_var, x_vec_var, context_vec_vars...; in_place=false)
     der2_exe! = make_function(der2_vec_var, x_vec_var, context_vec_vars...; in_place=true)
 
-    derivative_prep = DI.prepare_derivative(f, backend, x, contexts...)
+    derivative_prep = DI.prepare_derivative_nokwarg(strict, f, backend, x, contexts...)
     return FastDifferentiationAllocatingSecondDerivativePrep(
         _sig, y_prototype, derivative_prep, der2_exe, der2_exe!
     )
@@ -534,7 +534,7 @@ struct FastDifferentiationHVPPrep{SIG,E2,E2!,E1} <: DI.HVPPrep{SIG}
     gradient_prep::E1
 end
 
-function DI.prepare_hvp(
+function DI.prepare_hvp_nokwarg(
     strict::Val,
     f,
     backend::AutoFastDifferentiation,
@@ -557,7 +557,7 @@ function DI.prepare_hvp(
         hv_vec_var, x_vec_var, v_vec_var, context_vec_vars...; in_place=true
     )
 
-    gradient_prep = DI.prepare_gradient(f, backend, x, contexts...)
+    gradient_prep = DI.prepare_gradient_nokwarg(strict, f, backend, x, contexts...)
     return FastDifferentiationHVPPrep(_sig, hvp_exe, hvp_exe!, gradient_prep)
 end
 
@@ -633,7 +633,7 @@ struct FastDifferentiationHessianPrep{SIG,G,E2,E2!} <: DI.HessianPrep{SIG}
     hess_exe!::E2!
 end
 
-function DI.prepare_hessian(
+function DI.prepare_hessian_nokwarg(
     strict::Val,
     f,
     backend::Union{AutoFastDifferentiation,AutoSparse{<:AutoFastDifferentiation}},
@@ -656,7 +656,9 @@ function DI.prepare_hessian(
     hess_exe = make_function(hess_var, x_vec_var, context_vec_vars...; in_place=false)
     hess_exe! = make_function(hess_var, x_vec_var, context_vec_vars...; in_place=true)
 
-    gradient_prep = DI.prepare_gradient(f, dense_ad(backend), x, contexts...)
+    gradient_prep = DI.prepare_gradient_nokwarg(
+        strict, f, dense_ad(backend), x, contexts...
+    )
     return FastDifferentiationHessianPrep(_sig, gradient_prep, hess_exe, hess_exe!)
 end
 

DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/twoarg.jl

@@ -6,7 +6,7 @@ struct FastDifferentiationTwoArgPushforwardPrep{SIG,E1,E1!} <: DI.PushforwardPre
     jvp_exe!::E1!
 end
 
-function DI.prepare_pushforward(
+function DI.prepare_pushforward_nokwarg(
     strict::Val,
     f!,
     y,
@@ -107,7 +107,7 @@ struct FastDifferentiationTwoArgPullbackPrep{SIG,E1,E1!} <: DI.PullbackPrep{SIG}
     vjp_exe!::E1!
 end
 
-function DI.prepare_pullback(
+function DI.prepare_pullback_nokwarg(
     strict::Val,
     f!,
     y,
@@ -213,7 +213,7 @@ struct FastDifferentiationTwoArgDerivativePrep{SIG,E1,E1!} <: DI.DerivativePrep{
     der_exe!::E1!
 end
 
-function DI.prepare_derivative(
+function DI.prepare_derivative_nokwarg(
     strict::Val, f!, y, backend::AutoFastDifferentiation, x, contexts::Vararg{DI.Context,C};
 ) where {C}
     _sig = DI.signature(f!, y, backend, x, contexts...; strict)
@@ -295,7 +295,7 @@ struct FastDifferentiationTwoArgJacobianPrep{SIG,E1,E1!} <: DI.JacobianPrep{SIG}
     jac_exe!::E1!
 end
 
-function DI.prepare_jacobian(
+function DI.prepare_jacobian_nokwarg(
     strict::Val,
     f!,
     y,