perf: compute in-place HVP from in-place gradient

Author: gdalle

DifferentiationInterface/src/first_order/gradient.jl

@@ -132,3 +132,16 @@ function shuffled_gradient(
 ) where {F,C}
     return gradient(f, prep, backend, x, rewrap(unannotated_contexts...)...)
 end
+
+function shuffled_gradient!(
+    grad,
+    x,
+    f::F,
+    prep::GradientPrep,
+    backend::AbstractADType,
+    rewrap::Rewrap{C},
+    unannotated_contexts::Vararg{Any,C},
+) where {F,C}
+    gradient!(f, grad, prep, backend, x, rewrap(unannotated_contexts...)...)
+    return nothing
+end

DifferentiationInterface/src/misc/from_primitive.jl

@@ -14,10 +14,16 @@ Wrapper which forces a given backend to act as a reverse-mode backend.
 
 Used in internal testing.
 """
-struct AutoReverseFromPrimitive{B} <: FromPrimitive
+struct AutoReverseFromPrimitive{inplace,B<:AbstractADType} <: FromPrimitive
     backend::B
 end
 
+function AutoReverseFromPrimitive(backend::AbstractADType; inplace=false)
+    return AutoReverseFromPrimitive{inplace,typeof(backend)}(backend)
+end
+
+inplace_support(::AutoReverseFromPrimitive{true}) = InPlaceSupported()
+inplace_support(::AutoReverseFromPrimitive{false}) = InPlaceNotSupported()
 ADTypes.mode(::AutoReverseFromPrimitive) = ADTypes.ReverseMode()
 
 function threshold_batchsize(fromprim::AutoReverseFromPrimitive, dimension::Integer)