feat!: specify preparation arguments in DIT Scenario (#786)

* feat!: specify preparation arguments in DIT `Scenario`

* Fix

* Fixes

* Fixes

* Fixes

* Fix static arrays

* Fix

* Fix sparse and complex

* All works except HVP

* Fix tangents for prep same point

* Fixes

* Update DifferentiationInterfaceTest/src/scenarios/scenario.jl

Author: gdalle

DifferentiationInterface/src/utils/prep.jl

@@ -198,7 +198,7 @@ function check_prep(
         if SIG != EXEC_SIG
             throw(
                 PreparationMismatchError(
-                    SIG, EXEC_SIG; format=[:f, :backend, :x, :tang, :contexts]
+                    SIG, EXEC_SIG; format=[:f, :backend, :x, :t, :contexts]
                 ),
             )
         end
@@ -213,7 +213,7 @@ function check_prep(
         if SIG != EXEC_SIG
             throw(
                 PreparationMismatchError(
-                    SIG, EXEC_SIG; format=[:f!, :y, :backend, :x, :tang, :contexts]
+                    SIG, EXEC_SIG; format=[:f!, :y, :backend, :x, :t, :contexts]
                 ),
             )
         end

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

DifferentiationInterface/test/Back/FiniteDiff/test.jl

@@ -27,6 +27,7 @@ end
             include_cachified=true,
             include_constantorcachified=true,
             use_tuples=true,
+            include_smaller=true,
         );
         excluded=[:second_derivative, :hvp],
         logging=LOGGING,

DifferentiationInterface/test/Back/ForwardDiff/test.jl

@@ -41,6 +41,7 @@ end
             include_cachified=true,
             include_constantorcachified=true,
             use_tuples=true,
+            include_smaller=true,
         );
         logging=LOGGING,
     )

DifferentiationInterface/test/Core/Internals/signature.jl

@@ -98,7 +98,7 @@ end
         - exec: Nothing
       - backend: ✅
       - x: ✅
-      - tang: ✅
+      - t: ✅
       - contexts: ✅
     """ pushforward(nothing, prep, backend, x, (x,), Constant(c))
 end
@@ -119,7 +119,7 @@ end
       - y: ✅
       - backend: ✅
       - x: ✅
-      - tang: ✅
+      - t: ✅
       - contexts: ✅
     """ pushforward(nothing, y, prep, backend, x, (x,))
 end

DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl

@@ -63,7 +63,7 @@ end
 @testset "Dense" begin
     test_differentiation(
         vcat(backends, second_order_backends),
-        default_scenarios(; include_constantified=true);
+        default_scenarios(; include_constantified=true, include_smaller=true);
         logging=LOGGING,
     )
 

DifferentiationInterfaceTest/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Changed
+
+- Specify preparation arguments in DIT Scenario ([#786])
+
 ## [0.9.6] - 2025-03-28
 
 ### Added
@@ -18,6 +22,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 [unreleased]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterfaceTest-v0.9.6...main
 [0.9.6]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterfaceTest-v0.9.5...DifferentiationInterfaceTest-v0.9.6
 
+[#786]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/786
 [#749]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/749
 [#748]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/748
 [#745]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/745

DifferentiationInterfaceTest/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterfaceTest"
 uuid = "a82114a7-5aa3-49a8-9643-716bb13727a3"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.9.6"
+version = "0.10.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -44,7 +44,7 @@ AllocCheck = "0.2"
 Chairmarks = "1.2.1"
 ComponentArrays = "0.15"
 DataFrames = "1.6.1"
-DifferentiationInterface = "0.6.0"
+DifferentiationInterface = "0.6.53"
 DocStringExtensions = "0.8,0.9"
 ExplicitImports = "1.10.1"
 FiniteDiff = "2.27.0"

DifferentiationInterfaceTest/ext/DifferentiationInterfaceTestComponentArraysExt/DifferentiationInterfaceTestComponentArraysExt.jl

@@ -33,15 +33,13 @@ function comp_to_num_scenarios_onearg(x::ComponentVector; dx::AbstractVector, dy
         append!(
             scens,
             [
-                DIT.Scenario{:pullback,pl_op}(f, x; tang=(dy,), res1=(dx_from_dy,)),
+                DIT.Scenario{:pullback,pl_op}(f, x, (dy,); res1=(dx_from_dy,)),
                 DIT.Scenario{:gradient,pl_op}(f, x; res1=grad),
             ],
         )
     end
     for pl_op in (:out,)
-        append!(
-            scens, [DIT.Scenario{:pushforward,pl_op}(f, x; tang=(dx,), res1=(dy_from_dx,))]
-        )
+        append!(scens, [DIT.Scenario{:pushforward,pl_op}(f, x, (dx,); res1=(dy_from_dx,))])
     end
     return scens
 end

DifferentiationInterfaceTest/ext/DifferentiationInterfaceTestFluxExt/DifferentiationInterfaceTestFluxExt.jl

@@ -92,7 +92,11 @@ function DIT.flux_scenarios(rng::AbstractRNG=default_rng())
     g = gradient_finite_differences(square_loss, model, x)
 
     scen = DIT.Scenario{:gradient,:out}(
-        square_loss, model; contexts=(DI.Constant(x),), res1=g
+        square_loss,
+        model,
+        DI.Constant(x);
+        prep_args=(x=model, contexts=(DI.Constant(x),)),
+        res1=g,
     )
     push!(scens, scen)
 
@@ -163,7 +167,11 @@ function DIT.flux_scenarios(rng::AbstractRNG=default_rng())
         Flux.trainmode!(model)
         g = gradient_finite_differences(square_loss, model, x)
         scen = DIT.Scenario{:gradient,:out}(
-            square_loss, model; contexts=(DI.Constant(x),), res1=g
+            square_loss,
+            model,
+            DI.Constant(x);
+            prep_args=(; x=model, contexts=(DI.Constant(x),)),
+            res1=g,
         )
         push!(scens, scen)
     end
@@ -191,7 +199,11 @@ function DIT.flux_scenarios(rng::AbstractRNG=default_rng())
         Flux.trainmode!(model)
         g = gradient_finite_differences(square_loss_iterated, model, x)
         scen = DIT.Scenario{:gradient,:out}(
-            square_loss_iterated, model; contexts=(DI.Constant(x),), res1=g
+            square_loss_iterated,
+            model,
+            DI.Constant(x);
+            prep_args=(; x=model, contexts=(DI.Constant(x),)),
+            res1=g,
         )
         push!(scens, scen)
     end