Fixes

Author: gdalle

DifferentiationInterfaceTest/ext/DifferentiationInterfaceTestFluxExt/DifferentiationInterfaceTestFluxExt.jl

@@ -162,7 +162,13 @@ function DIT.flux_scenarios(rng::AbstractRNG=default_rng())
     for (model, x) in models_and_xs
         Flux.trainmode!(model)
         g = gradient_finite_differences(square_loss, model, x)
-        scen = DIT.Scenario{:gradient,:out}(square_loss, model, DI.Constant(x); res1=g)
+        scen = DIT.Scenario{:gradient,:out}(
+            square_loss,
+            model,
+            DI.Constant(x);
+            prep_args=(; x=model, contexts=(DI.Constant(x),)),
+            res1=g,
+        )
         push!(scens, scen)
     end
 
@@ -189,7 +195,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, 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

DifferentiationInterfaceTest/ext/DifferentiationInterfaceTestJLArraysExt/DifferentiationInterfaceTestJLArraysExt.jl

@@ -23,13 +23,14 @@ myjl(x::DI.Cache{<:Union{Tuple,NamedTuple}}) = map(myjl, map(DI.Cache, DI.unwrap
 myjl(::Nothing) = nothing
 
 function myjl(scen::DIT.Scenario{op,pl_op,pl_fun}) where {op,pl_op,pl_fun}
-    (; f, x, y, t, contexts, res1, res2) = scen
+    (; f, x, y, t, contexts, prep_args, res1, res2) = scen
     return DIT.Scenario{op,pl_op,pl_fun}(;
         f=myjl(f),
         x=myjl(x),
         y=myjl(y),
         t=myjl(t),
         contexts=myjl(contexts),
+        prep_args=map(myjl, prep_args),
         res1=myjl(res1),
         res2=myjl(res2),
     )

DifferentiationInterfaceTest/ext/DifferentiationInterfaceTestLuxExt/DifferentiationInterfaceTestLuxExt.jl

@@ -203,6 +203,10 @@ function DIT.lux_scenarios(rng::AbstractRNG=default_rng())
             DI.Constant(model),
             DI.Constant(x),
             DI.Constant(st);
+            prep_args=(
+                x=ComponentArray(ps),
+                contexts=(DI.Constant(model), DI.Constant(x), DI.Constant(st)),
+            ),
             res1=g,
         )
         push!(scens, scen)

DifferentiationInterfaceTest/ext/DifferentiationInterfaceTestStaticArraysExt/DifferentiationInterfaceTestStaticArraysExt.jl

@@ -36,13 +36,14 @@ end
 mystatic(::Nothing) = nothing
 
 function mystatic(scen::DIT.Scenario{op,pl_op,pl_fun}) where {op,pl_op,pl_fun}
-    (; f, x, y, t, contexts, res1, res2) = scen
+    (; f, x, y, t, contexts, prep_args, res1, res2) = scen
     return DIT.Scenario{op,pl_op,pl_fun}(;
         f=mystatic(f),
         x=mystatic(x),
         y=pl_fun == :in ? mymutablestatic(y) : mystatic(y),
         t=mystatic(t),
         contexts=mystatic(contexts),
+        prep_args=map(mystatic, prep_args),
         res1=mystatic(res1),
         res2=mystatic(res2),
     )

DifferentiationInterfaceTest/src/utils.jl

@@ -6,14 +6,6 @@ myzero(::Nothing) = nothing
 mysimilar(x::Number) = one(x)
 mysimilar(x::AbstractArray) = similar(x)
 mysimilar(x::Union{Tuple,NamedTuple}) = map(mysimilar, x)
-mysimilar(x) = deepcopy(x)
-
-myrandom(rng::AbstractRNG, x::Number) = randn(rng, typeof(x))
-myrandom(rng::AbstractRNG, x::AbstractArray) = map(Base.Fix1(myrandom, rng), x)
-myrandom(rng::AbstractRNG, x::Union{Tuple,NamedTuple}) = map(Base.Fix1(myrandom, rng), x)
-myrandom(rng::AbstractRNG, x) = deepcopy(x)
-
-myrandom(x) = myrandom(default_rng(), x)
 
 mysize(x::Number) = size(x)
 mysize(x::AbstractArray) = size(x)