test: allow passing other scenarios to adjust preparation from (#701)

* test: allow passing other scenarios to adjust preparation from

* Typo

* copy

* Switch to other scenario bundled with initial scenario

* No fail fast

* Typo

* better batchsize handling

* Fix

* fiox

* fix

* fix scenarios

* Smaller scens

* Fix tests

* Fix tracker

* Use with_contexts for Tracker

* fixes

* Revert Tracker changes

Author: gdalle

.github/workflows/Test.yml

@@ -25,7 +25,7 @@ jobs:
       actions: write
       contents: read
     strategy:
-      fail-fast: true  # TODO: toggle
+      fail-fast: false  # TODO: toggle
       matrix:
         version:
           - "1.10"

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.6.34"
+version = "0.6.35"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl

@@ -1,5 +1,5 @@
 # until https://github.com/EnzymeAD/Enzyme.jl/pull/1545 is merged
-function DI.BatchSizeSettings(::AutoEnzyme, N::Integer)
+function DI.pick_batchsize(::AutoEnzyme, N::Integer)
     B = DI.reasonable_batchsize(N, 16)
     return DI.BatchSizeSettings{B}(N)
 end

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl

@@ -199,6 +199,23 @@ function DI.prepare_derivative(
     return ForwardDiffTwoArgDerivativePrep(config)
 end
 
+function DI.prepare!_derivative(
+    f!::F,
+    y,
+    old_prep::ForwardDiffTwoArgDerivativePrep,
+    backend::AutoForwardDiff,
+    x,
+    contexts::Vararg{DI.ConstantOrFunctionOrBackend,C},
+) where {F,C}
+    if y isa Vector
+        (; config) = old_prep
+        resize!(config.duals, length(y))
+        return old_prep
+    else
+        return DI.prepare_derivative(f!, y, backend, x, contexts...)
+    end
+end
+
 function DI.value_and_derivative(
     f!::F,
     y,
@@ -352,6 +369,25 @@ function DI.prepare_jacobian(
     return ForwardDiffTwoArgJacobianPrep(config)
 end
 
+function DI.prepare!_jacobian(
+    f!::F,
+    y,
+    old_prep::ForwardDiffTwoArgJacobianPrep,
+    backend::AutoForwardDiff,
+    x,
+    contexts::Vararg{DI.ConstantOrFunctionOrBackend,C},
+) where {F,C}
+    if x isa Vector && y isa Vector
+        (; config) = old_prep
+        (yduals, xduals) = config.duals
+        resize!(yduals, length(y))
+        resize!(xduals, length(x))
+        return old_prep
+    else
+        return DI.prepare_jacobian(f!, y, backend, x, contexts...)
+    end
+end
+
 function DI.value_and_jacobian(
     f!::F,
     y,

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl

@@ -1,9 +1,9 @@
-function DI.BatchSizeSettings(::AutoForwardDiff{nothing}, N::Integer)
+function DI.pick_batchsize(::AutoForwardDiff{nothing}, N::Integer)
     chunksize = ForwardDiff.pickchunksize(N)
     return DI.BatchSizeSettings{chunksize}(N)
 end
 
-function DI.BatchSizeSettings(::AutoForwardDiff{chunksize}, N::Integer) where {chunksize}
+function DI.pick_batchsize(::AutoForwardDiff{chunksize}, N::Integer) where {chunksize}
     return DI.BatchSizeSettings{chunksize}(N)
 end
 

DifferentiationInterface/ext/DifferentiationInterfacePolyesterForwardDiffExt/DifferentiationInterfacePolyesterForwardDiffExt.jl

@@ -13,12 +13,12 @@ end
 
 DI.check_available(::AutoPolyesterForwardDiff) = true
 
-function DI.BatchSizeSettings(backend::AutoPolyesterForwardDiff, x::AbstractArray)
-    return DI.BatchSizeSettings(single_threaded(backend), x)
+function DI.pick_batchsize(backend::AutoPolyesterForwardDiff, x::AbstractArray)
+    return DI.pick_batchsize(single_threaded(backend), x)
 end
 
-function DI.BatchSizeSettings(backend::AutoPolyesterForwardDiff, N::Integer)
-    return DI.BatchSizeSettings(single_threaded(backend), N)
+function DI.pick_batchsize(backend::AutoPolyesterForwardDiff, N::Integer)
+    return DI.pick_batchsize(single_threaded(backend), N)
 end
 
 function DI.threshold_batchsize(

DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl

@@ -14,27 +14,25 @@ end
 
 DI.ismutable_array(::Type{<:SArray}) = false
 
-function DI.BatchSizeSettings(::DI.AutoSimpleFiniteDiff{nothing}, x::StaticArray)
+function DI.pick_batchsize(::DI.AutoSimpleFiniteDiff{nothing}, x::StaticArray)
     return DI.BatchSizeSettings{length(x),true,true}(length(x))
 end
 
-function DI.BatchSizeSettings(::AutoForwardDiff{nothing}, x::StaticArray)
+function DI.pick_batchsize(::AutoForwardDiff{nothing}, x::StaticArray)
     return DI.BatchSizeSettings{length(x),true,true}(length(x))
 end
 
-function DI.BatchSizeSettings(::AutoEnzyme, x::StaticArray)
+function DI.pick_batchsize(::AutoEnzyme, x::StaticArray)
     return DI.BatchSizeSettings{length(x),true,true}(length(x))
 end
 
-function DI.BatchSizeSettings(
+function DI.pick_batchsize(
     ::DI.AutoSimpleFiniteDiff{chunksize}, x::StaticArray
 ) where {chunksize}
     return DI.BatchSizeSettings{chunksize}(Val(length(x)))
 end
 
-function DI.BatchSizeSettings(
-    ::AutoForwardDiff{chunksize}, x::StaticArray
-) where {chunksize}
+function DI.pick_batchsize(::AutoForwardDiff{chunksize}, x::StaticArray) where {chunksize}
     return DI.BatchSizeSettings{chunksize}(Val(length(x)))
 end
 

DifferentiationInterface/src/fallbacks/change_prep.jl

@@ -43,22 +43,22 @@ for op in [
     if op in (:derivative, :gradient, :jacobian)
         # 1-arg
         @eval function $prep_op!(
-            f::F, ::$P, backend::AbstractADType, x, contexts::Vararg{Context,C}
+            f::F, old_prep::$P, backend::AbstractADType, x, contexts::Vararg{Context,C}
         ) where {F,C}
             return $prep_op(f, backend, x, contexts...)
         end
         op == :gradient && continue
         # 2-arg
         @eval function $prep_op!(
-            f!::F, y, ::$P, backend::AbstractADType, x, contexts::Vararg{Context,C}
+            f!::F, y, old_prep::$P, backend::AbstractADType, x, contexts::Vararg{Context,C}
         ) where {F,C}
             return $prep_op(f!, y, backend, x, contexts...)
         end
 
     elseif op in (:second_derivative, :hessian)
         # 1-arg
         @eval function $prep_op!(
-            f::F, ::$P, backend::AbstractADType, x, contexts::Vararg{Context,C}
+            f::F, old_prep::$P, backend::AbstractADType, x, contexts::Vararg{Context,C}
         ) where {F,C}
             return $prep_op(f, backend, x, contexts...)
         end
@@ -67,7 +67,7 @@ for op in [
         # 1-arg
         @eval function $prep_op!(
             f::F,
-            ::$P,
+            old_prep::$P,
             backend::AbstractADType,
             x,
             seed::NTuple,
@@ -96,7 +96,7 @@ for op in [
         @eval function $prep_op!(
             f!::F,
             y,
-            ::$P,
+            old_prep::$P,
             backend::AbstractADType,
             x,
             seed::NTuple,

DifferentiationInterface/src/first_order/jacobian.jl

@@ -100,9 +100,9 @@ function prepare_jacobian(
     perf = pushforward_performance(backend)
     # type-unstable
     if perf isa PushforwardFast
-        batch_size_settings = BatchSizeSettings(backend, x)
+        batch_size_settings = pick_batchsize(backend, x)
     else
-        batch_size_settings = BatchSizeSettings(backend, y)
+        batch_size_settings = pick_batchsize(backend, y)
     end
     # function barrier
     return _prepare_jacobian_aux(
@@ -116,9 +116,9 @@ function prepare_jacobian(
     perf = pushforward_performance(backend)
     # type-unstable
     if perf isa PushforwardFast
-        batch_size_settings = BatchSizeSettings(backend, x)
+        batch_size_settings = pick_batchsize(backend, x)
     else
-        batch_size_settings = BatchSizeSettings(backend, y)
+        batch_size_settings = pick_batchsize(backend, y)
     end
     # function barrier
     return _prepare_jacobian_aux(

DifferentiationInterface/src/misc/from_primitive.jl

@@ -3,8 +3,8 @@ abstract type FromPrimitive <: AbstractADType end
 check_available(fromprim::FromPrimitive) = check_available(fromprim.backend)
 inplace_support(fromprim::FromPrimitive) = inplace_support(fromprim.backend)
 
-function pick_batchsize(fromprim::FromPrimitive, x_or_N::Union{AbstractArray,Integer})
-    return pick_batchsize(fromprim.backend, x_or_N)
+function pick_batchsize(fromprim::FromPrimitive, N::Integer)
+    return pick_batchsize(fromprim.backend, N)
 end
 
 struct AutoReverseFromPrimitive{B} <: FromPrimitive