Last fix

Author: gdalle

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(f, backend, x, (y,); strict=Val(true))
+    prep_same = DI.prepare_pullback_same_point(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

@@ -19,7 +19,6 @@ function DI.prepare_pullback(
 end
 
 function DI.prepare_pullback_same_point(
-    strict,
     f,
     prep::DI.NoPullbackPrep,
     backend::AutoReverseChainRules,
@@ -28,7 +27,7 @@ function DI.prepare_pullback_same_point(
     contexts::Vararg{DI.GeneralizedConstant,C};
 ) where {C}
     DI.check_prep(f, prep, backend, x, ty, contexts...)
-    _sig = DI.signature(f, backend, x, ty, contexts...; strict)
+    _sig = DI.signature(f, backend, x, ty, contexts...; strict=DI.is_strict(prep))
     rc = ruleconfig(backend)
     y, pb = rrule_via_ad(rc, f, x, map(DI.unwrap, contexts)...)
     return ChainRulesPullbackPrepSamePoint(_sig, y, pb)

DifferentiationInterface/ext/DifferentiationInterfaceTrackerExt/DifferentiationInterfaceTrackerExt.jl

@@ -35,8 +35,8 @@ function DI.prepare_pullback_same_point(
     ty::NTuple,
     contexts::Vararg{DI.GeneralizedConstant,C},
 ) where {C}
-    _sig = DI.signature(f, backend, x, ty, contexts...; strict=DI.is_strict(prep))
     DI.check_prep(f, prep, backend, x, ty, contexts...)
+    _sig = DI.signature(f, backend, x, ty, contexts...; strict=DI.is_strict(prep))
     y, pb = forward(f, x, map(DI.unwrap, contexts)...)
     return TrackerPullbackPrepSamePoint(_sig, y, pb)
 end

DifferentiationInterface/src/docstrings.jl

@@ -19,15 +19,14 @@ function docstring_prepare(operator; samepoint=false, inplace=false)
     Create a `prep` object that can be given to [`$(operator)`](@ref) and its variants to speed them up$(samepoint_warning(samepoint)).
 
     Depending on the backend, this can have several effects (preallocating memory, recording an execution trace) which are transparent to the user.
+    $(inplace ? "\nFor in-place functions, `y` is mutated by `f!` during preparation." : "")
 
     !!! warning
         The preparation result `prep` is only reusable as long as the arguments to `$operator` do not change type or size, and the function and backend themselves are not modified.
         Otherwise, preparation becomes invalid and you need to run it again.
         In some settings, invalid preparations may still give correct results (e.g. for backends that require no preparation), but this is not a semantic guarantee and should not be relied upon.
 
     When `strict=Val(true)`, type checking is enforced between preparation and execution (but size checking is left to the user).
-
-    $(inplace ? "\nFor in-place functions, `y` is mutated by `f!` during preparation." : "")
     """
 end