Switch order of arguments in preparation (#177)

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_twoarg.jl

@@ -1,6 +1,6 @@
 ## Pushforward
 
-DI.prepare_pushforward(f!, ::AutoForwardEnzyme, y, x) = NoPushforwardExtras()
+DI.prepare_pushforward(f!, y, ::AutoForwardEnzyme, x) = NoPushforwardExtras()
 
 function DI.value_and_pushforward(
     f!, y, backend::AutoForwardEnzyme, x, dx, ::NoPushforwardExtras

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_twoarg.jl

@@ -1,6 +1,6 @@
 ## Pullback
 
-DI.prepare_pullback(f!, ::AutoReverseEnzyme, y, x) = NoPullbackExtras()
+DI.prepare_pullback(f!, y, ::AutoReverseEnzyme, x) = NoPullbackExtras()
 
 function DI.value_and_pullback(
     f!, y, ::AutoReverseEnzyme, x::Number, dy, ::NoPullbackExtras

DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/twoarg.jl

@@ -5,7 +5,7 @@ struct FastDifferentiationTwoArgPushforwardExtras{E1,E2} <: PushforwardExtras
     jvp_exe!::E2
 end
 
-function DI.prepare_pushforward(f!, ::AnyAutoFastDifferentiation, y, x)
+function DI.prepare_pushforward(f!, y, ::AnyAutoFastDifferentiation, x)
     x_var = if x isa Number
         only(make_variables(:x))
     else
@@ -85,7 +85,7 @@ struct FastDifferentiationTwoArgDerivativeExtras{E1,E2} <: DerivativeExtras
     der_exe!::E2
 end
 
-function DI.prepare_derivative(f!, ::AnyAutoFastDifferentiation, y, x)
+function DI.prepare_derivative(f!, y, ::AnyAutoFastDifferentiation, x)
     x_var = only(make_variables(:x))
     y_var = make_variables(:y, size(y)...)
     f!(y_var, x_var)
@@ -153,7 +153,7 @@ struct FastDifferentiationTwoArgJacobianExtras{E1,E2} <: JacobianExtras
     jac_exe!::E2
 end
 
-function DI.prepare_jacobian(f!, backend::AnyAutoFastDifferentiation, y, x)
+function DI.prepare_jacobian(f!, y, backend::AnyAutoFastDifferentiation, x)
     x_var = make_variables(:x, size(x)...)
     y_var = make_variables(:y, size(y)...)
     f!(y_var, x_var)

DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/twoarg.jl

@@ -1,6 +1,6 @@
 ## Pushforward
 
-DI.prepare_pushforward(f!, ::AnyAutoFiniteDiff, y, x) = NoPushforwardExtras()
+DI.prepare_pushforward(f!, y, ::AnyAutoFiniteDiff, x) = NoPushforwardExtras()
 
 function DI.value_and_pushforward(
     f!, y, backend::AnyAutoFiniteDiff, x, dx, ::NoPushforwardExtras
@@ -23,7 +23,7 @@ struct FiniteDiffTwoArgDerivativeExtras{C}
     cache::C
 end
 
-function DI.prepare_derivative(f!, ::AnyAutoFiniteDiff, y, x)
+function DI.prepare_derivative(f!, y, ::AnyAutoFiniteDiff, x)
     cache = nothing
     return FiniteDiffTwoArgDerivativeExtras(cache)
 end
@@ -65,7 +65,7 @@ struct FiniteDiffTwoArgJacobianExtras{C}
     cache::C
 end
 
-function DI.prepare_jacobian(f!, backend::AnyAutoFiniteDiff, y, x)
+function DI.prepare_jacobian(f!, y, backend::AnyAutoFiniteDiff, x)
     x1 = similar(x)
     fx = similar(y)
     fx1 = similar(y)

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl

@@ -1,4 +1,4 @@
-DI.prepare_pushforward(f!, ::AnyAutoForwardDiff, y, x) = NoPushforwardExtras()
+DI.prepare_pushforward(f!, y, ::AnyAutoForwardDiff, x) = NoPushforwardExtras()
 
 function DI.value_and_pushforward(f!, y, ::AnyAutoForwardDiff, x, dx, ::NoPushforwardExtras)
     T = tag_type(f!, x)
@@ -16,7 +16,7 @@ struct ForwardDiffTwoArgDerivativeExtras{C} <: DerivativeExtras
     config::C
 end
 
-function DI.prepare_derivative(f!, ::AnyAutoForwardDiff, y::AbstractArray, x::Number)
+function DI.prepare_derivative(f!, y::AbstractArray, ::AnyAutoForwardDiff, x::Number)
     return ForwardDiffTwoArgDerivativeExtras(DerivativeConfig(f!, y, x))
 end
 
@@ -75,7 +75,7 @@ struct ForwardDiffTwoArgJacobianExtras{C} <: JacobianExtras
 end
 
 function DI.prepare_jacobian(
-    f!, backend::AnyAutoForwardDiff, y::AbstractArray, x::AbstractArray
+    f!, y::AbstractArray, backend::AnyAutoForwardDiff, x::AbstractArray
 )
     return ForwardDiffTwoArgJacobianExtras(
         JacobianConfig(f!, y, x, choose_chunk(backend, x))

DifferentiationInterface/ext/DifferentiationInterfacePolyesterForwardDiffExt/twoarg.jl

@@ -1,7 +1,7 @@
 ## Pushforward
 
-function DI.prepare_pushforward(f!, backend::AnyAutoPolyForwardDiff, y, x)
-    return DI.prepare_pushforward(f!, single_threaded(backend), y, x)
+function DI.prepare_pushforward(f!, y, backend::AnyAutoPolyForwardDiff, x)
+    return DI.prepare_pushforward(f!, y, single_threaded(backend), x)
 end
 
 function DI.value_and_pushforward(
@@ -30,8 +30,8 @@ end
 
 ## Derivative
 
-function DI.prepare_derivative(f!, backend::AnyAutoPolyForwardDiff, y, x)
-    return DI.prepare_derivative(f!, single_threaded(backend), y, x)
+function DI.prepare_derivative(f!, y, backend::AnyAutoPolyForwardDiff, x)
+    return DI.prepare_derivative(f!, y, single_threaded(backend), x)
 end
 
 function DI.value_and_derivative(
@@ -58,7 +58,7 @@ end
 
 ## Jacobian
 
-DI.prepare_jacobian(f!, ::AnyAutoPolyForwardDiff, y, x) = NoJacobianExtras()
+DI.prepare_jacobian(f!, y, ::AnyAutoPolyForwardDiff, x) = NoJacobianExtras()
 
 function DI.value_and_jacobian(
     f!, y, ::AnyAutoPolyForwardDiff{C}, x, ::NoJacobianExtras

DifferentiationInterface/ext/DifferentiationInterfaceReverseDiffExt/twoarg.jl

@@ -1,6 +1,6 @@
 ## Pullback
 
-DI.prepare_pullback(f!, ::AnyAutoReverseDiff, y, x) = NoPullbackExtras()
+DI.prepare_pullback(f!, y, ::AnyAutoReverseDiff, x) = NoPullbackExtras()
 
 ### Array in
 
@@ -60,7 +60,7 @@ function DI.value_and_pullback(
     x_array = [x]
     dx_array = similar(x_array)
     f!_array(_y::AbstractArray, _x_array) = f!(_y, only(_x_array))
-    new_extras = DI.prepare_pullback(f!_array, backend, y, x_array)
+    new_extras = DI.prepare_pullback(f!_array, y, backend, x_array)
     y, dx_array = DI.value_and_pullback(f!_array, y, backend, x_array, dy, new_extras)
     return y, only(dx_array)
 end
@@ -72,7 +72,7 @@ struct ReverseDiffTwoArgJacobianExtras{T} <: JacobianExtras
 end
 
 function DI.prepare_jacobian(
-    f!, backend::AnyAutoReverseDiff, y::AbstractArray, x::AbstractArray
+    f!, y::AbstractArray, backend::AnyAutoReverseDiff, x::AbstractArray
 )
     tape = JacobianTape(f!, y, x)
     if backend.compile

DifferentiationInterface/ext/DifferentiationInterfaceSparseDiffToolsExt/twoarg.jl

@@ -7,7 +7,7 @@ for AutoSparse in SPARSE_BACKENDS
         ## Jacobian
 
         function DI.prepare_jacobian(
-            f!, backend::$AutoSparse, y::AbstractArray, x::AbstractArray
+            f!, y::AbstractArray, backend::$AutoSparse, x::AbstractArray
         )
             cache = sparse_jacobian_cache(
                 backend, SymbolicsSparsityDetection(), f!, similar(y), x

DifferentiationInterface/ext/DifferentiationInterfaceTapirExt/twoarg.jl

@@ -2,7 +2,7 @@ struct TapirTwoArgPullbackExtras{R} <: PullbackExtras
     rrule::R
 end
 
-function DI.prepare_pullback(f!, ::AutoTapir, y, x)
+function DI.prepare_pullback(f!, y, ::AutoTapir, x)
     return TapirTwoArgPullbackExtras(build_rrule(f!, y, x))
 end
 

DifferentiationInterface/src/derivative.jl

@@ -1,10 +1,12 @@
 ## Docstrings
 
 """
-    prepare_derivative(f,  backend,    x) -> extras
-    prepare_derivative(f!, backend, y, x) -> extras
+    prepare_derivative(f,     backend, x) -> extras
+    prepare_derivative(f!, y, backend, x) -> extras
 
 Create an `extras` object subtyping [`DerivativeExtras`](@ref) that can be given to derivative operators.
+
+Beware that in the two-argument case, `y` is mutated by `f!` during preparation.
 """
 function prepare_derivative end
 
@@ -51,8 +53,8 @@ function prepare_derivative(f, backend::AbstractADType, x)
     return PushforwardDerivativeExtras(prepare_pushforward(f, backend, x))
 end
 
-function prepare_derivative(f!, backend::AbstractADType, y, x)
-    return PushforwardDerivativeExtras(prepare_pushforward(f!, backend, y, x))
+function prepare_derivative(f!, y, backend::AbstractADType, x)
+    return PushforwardDerivativeExtras(prepare_pushforward(f!, y, backend, x))
 end
 
 ## One argument
@@ -102,7 +104,7 @@ function value_and_derivative(
     y,
     backend::AbstractADType,
     x,
-    extras::DerivativeExtras=prepare_derivative(f!, backend, y, x),
+    extras::DerivativeExtras=prepare_derivative(f!, y, backend, x),
 )
     return value_and_pushforward(f!, y, backend, x, one(x), extras.pushforward_extras)
 end
@@ -113,7 +115,7 @@ function value_and_derivative!(
     der,
     backend::AbstractADType,
     x,
-    extras::DerivativeExtras=prepare_derivative(f!, backend, y, x),
+    extras::DerivativeExtras=prepare_derivative(f!, y, backend, x),
 )
     return value_and_pushforward!(f!, y, der, backend, x, one(x), extras.pushforward_extras)
 end
@@ -123,7 +125,7 @@ function derivative(
     y,
     backend::AbstractADType,
     x,
-    extras::DerivativeExtras=prepare_derivative(f!, backend, y, x),
+    extras::DerivativeExtras=prepare_derivative(f!, y, backend, x),
 )
     return pushforward(f!, y, backend, x, one(x), extras.pushforward_extras)
 end
@@ -134,7 +136,7 @@ function derivative!(
     der,
     backend::AbstractADType,
     x,
-    extras::DerivativeExtras=prepare_derivative(f!, backend, y, x),
+    extras::DerivativeExtras=prepare_derivative(f!, y, backend, x),
 )
     return pushforward!(f!, y, der, backend, x, one(x), extras.pushforward_extras)
 end