Fixes

Author: gdalle

DifferentiationInterface/test/Back/DifferentiateWith/test.jl

@@ -16,9 +16,7 @@ function differentiatewith_scenarios()
             DIT.function_place(scen) == :out
         end
     good_scens = map(bad_scens) do scen
-        DIT.change_function(
-            scen, DifferentiateWith(scen.f, AutoFiniteDiff()); keep_smaller=false
-        )
+        DIT.change_function(scen, DifferentiateWith(scen.f, AutoFiniteDiff()))
     end
     return good_scens
 end

DifferentiationInterfaceTest/src/scenarios/modify.jl

@@ -36,9 +36,7 @@ end
 
 Return a new `Scenario` identical to `scen` except for the function `f` which is changed to `new_f`.
 """
-function change_function(
-    scen::Scenario{op,pl_op,pl_fun}, new_f; keep_smaller
-) where {op,pl_op,pl_fun}
+function change_function(scen::Scenario{op,pl_op,pl_fun}, new_f) where {op,pl_op,pl_fun}
     return Scenario{op,pl_op,pl_fun}(;
         f=new_f,
         x=scen.x,
@@ -52,6 +50,8 @@ function change_function(
     )
 end
 
+same_function(scen) = change_function(scen, scen.f)
+
 """
     batchify(scen::Scenario)
 

DifferentiationInterfaceTest/src/scenarios/scenario.jl

@@ -199,13 +199,15 @@ function Base.show(
 end
 
 function adapt_batchsize(backend::AbstractADType, scen::Scenario)
-    (; x, y) = scen
+    (; x, y, prep_args) = scen
+    xprep = prep_args.x
+    yprep = hasproperty(prep_args, :y) ? prep_args.y : y
     Bmax = if x isa AbstractArray && y isa AbstractArray
-        min(length(x), length(y))
+        min(length(x), length(y), length(xprep), length(yprep))
     elseif x isa AbstractArray
-        length(x)
+        min(length(x), length(xprep))
     elseif y isa AbstractArray
-        length(y)
+        min(length(y), length(yprep))
     else
         typemax(Int)
     end

DifferentiationInterfaceTest/src/test_differentiation.jl

@@ -48,6 +48,7 @@ Each setting tests/benchmarks a different subset of calls:
 - `rtol=1e-3`: relative precision for correctness testing (when comparing to the reference outputs)
 - `scenario_intact=true`: whether to check that the scenario remains unchanged after the operators are applied
 - `sparsity=false`: whether to check sparsity patterns for Jacobians / Hessians
+- `reprepare::Bool=true`: whether to modify preparation before testing when the preparation arguments have the wrong size
 
 **Type stability options:**
 

DifferentiationInterfaceTest/test/standard.jl

@@ -20,7 +20,10 @@ test_differentiation(
 test_differentiation(
     [AutoForwardDiff(), AutoFiniteDiff(; relstep=1e-5)],
     default_scenarios(;
-        include_batchified=false, include_normal=false, include_constantorcachified=true
+        include_batchified=false,
+        include_normal=false,
+        include_cachified=true,
+        include_constantorcachified=true,
     );
     logging=LOGGING,
 )
@@ -35,7 +38,7 @@ sparse_backend = AutoSparse(
 
 test_differentiation(
     sparse_backend,
-    sparse_scenarios(; include_cachified=true, use_tuples=true);
+    sparse_scenarios(; include_cachified=true, use_tuples=false);
     sparsity=true,
     logging=LOGGING,
 )

DifferentiationInterfaceTest/test/zero_backends.jl

@@ -11,15 +11,17 @@ LOGGING = get(ENV, "CI", "false") == "false"
 
 test_differentiation(
     AutoZeroForward(),
-    default_scenarios(; include_batchified=false);
-    correctness=false,
+    map(zero, default_scenarios(; include_batchified=false));
     type_stability=:full,
     logging=LOGGING,
 )
 
 test_differentiation(
     AutoZeroReverse(),
-    default_scenarios(; include_batchified=false);
+    map(
+        DifferentiationInterfaceTest.same_function,
+        default_scenarios(; include_batchified=false),
+    );
     correctness=false,
     type_stability=:prepared,
     logging=LOGGING,