More sophisticated testing

gdalle · gdalle · commit 7f63ad91eec6 · 2025-03-18T15:17:43.000+01:00
diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/DifferentiationInterfaceFastDifferentiationExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/DifferentiationInterfaceFastDifferentiationExt.jl
@@ -23,11 +23,6 @@ myvec(x::AbstractArray) = vec(x)
 
 variablize(::Number, name::Symbol) = only(make_variables(name))
 variablize(x::AbstractArray, name::Symbol) = make_variables(name, size(x)...)
-function variablize(x::Union{Tuple,NamedTuple}, name::Symbol)
-    return map(x) do xi
-        variablize(xi, gensym())  # TODO: fix symbol?
-    end
-end
 
 function variablize(contexts::NTuple{C,DI.Context}) where {C}
     map(enumerate(contexts)) do (k, c)
diff --git a/DifferentiationInterface/src/utils/context.jl b/DifferentiationInterface/src/utils/context.jl
@@ -77,7 +77,7 @@ The initial values present inside the cache do not matter.
 For some backends, preparation allocates the required memory for `Cache` contexts with the right element type, similar to [PreallocationTools.jl](https://github.com/SciML/PreallocationTools.jl).
 
 !!! warning
-    Some backends require any `Cache` context to be an `AbstractArray or a (named) tuple of `AbstractArray`s.
+    Some backends require any `Cache` context to be an `AbstractArray`, others accept nested (named) tuples of `AbstractArray`s.
 
 # Example
 
diff --git a/DifferentiationInterface/test/Back/Enzyme/test.jl b/DifferentiationInterface/test/Back/Enzyme/test.jl
@@ -55,7 +55,7 @@ end;
 
     test_differentiation(
         backends[2],
-        default_scenarios(; include_normal=false, include_cachified=true);
+        default_scenarios(; include_normal=false, include_cachified=true, use_tuples=true);
         excluded=SECOND_ORDER,
         logging=LOGGING,
     )
diff --git a/DifferentiationInterface/test/Back/FiniteDiff/test.jl b/DifferentiationInterface/test/Back/FiniteDiff/test.jl
@@ -22,7 +22,9 @@ end
 @testset "Dense" begin
     test_differentiation(
         AutoFiniteDiff(),
-        default_scenarios(; include_constantified=true, include_cachified=true);
+        default_scenarios(;
+            include_constantified=true, include_cachified=true, use_tuples=true
+        );
         excluded=[:second_derivative, :hvp],
         logging=LOGGING,
     )
diff --git a/DifferentiationInterface/test/Back/FiniteDifferences/test.jl b/DifferentiationInterface/test/Back/FiniteDifferences/test.jl
@@ -19,7 +19,9 @@ end
 
 test_differentiation(
     AutoFiniteDifferences(; fdm=FiniteDifferences.central_fdm(3, 1)),
-    default_scenarios(; include_constantified=true, include_cachified=true);
+    default_scenarios(;
+        include_constantified=true, include_cachified=true, use_tuples=true
+    );
     excluded=SECOND_ORDER,
     logging=LOGGING,
 );
diff --git a/DifferentiationInterface/test/Back/ForwardDiff/test.jl b/DifferentiationInterface/test/Back/ForwardDiff/test.jl
@@ -36,7 +36,10 @@ end
     test_differentiation(
         AutoForwardDiff(),
         default_scenarios(;
-            include_normal=false, include_batchified=false, include_cachified=true
+            include_normal=false,
+            include_batchified=false,
+            include_cachified=true,
+            use_tuples=true,
         );
         logging=LOGGING,
     )
diff --git a/DifferentiationInterface/test/Back/Mooncake/test.jl b/DifferentiationInterface/test/Back/Mooncake/test.jl
@@ -19,7 +19,9 @@ end
 
 test_differentiation(
     backends,
-    default_scenarios(; include_constantified=true, include_cachified=true);
+    default_scenarios(;
+        include_constantified=true, include_cachified=true, use_tuples=true
+    );
     excluded=SECOND_ORDER,
     logging=LOGGING,
 );
diff --git a/DifferentiationInterface/test/Back/PolyesterForwardDiff/test.jl b/DifferentiationInterface/test/Back/PolyesterForwardDiff/test.jl
@@ -28,7 +28,9 @@ end
 
 test_differentiation(
     backends,
-    default_scenarios(; include_constantified=true, include_cachified=true);
+    default_scenarios(;
+        include_constantified=true, include_cachified=true, use_tuples=true
+    );
     logging=LOGGING,
 );
 
diff --git a/DifferentiationInterface/test/Back/SymbolicBackends/fastdifferentiation.jl b/DifferentiationInterface/test/Back/SymbolicBackends/fastdifferentiation.jl
@@ -17,7 +17,9 @@ end
 
 test_differentiation(
     AutoFastDifferentiation(),
-    default_scenarios(; include_constantified=true, include_cachified=true);
+    default_scenarios(;
+        include_constantified=true, include_cachified=true, use_tuples=false
+    );
     logging=LOGGING,
 );
 
diff --git a/DifferentiationInterface/test/Back/SymbolicBackends/symbolics.jl b/DifferentiationInterface/test/Back/SymbolicBackends/symbolics.jl
@@ -21,7 +21,7 @@ test_differentiation(
 
 test_differentiation(
     AutoSymbolics(),
-    default_scenarios(; include_normal=false, include_cachified=true);
+    default_scenarios(; include_normal=false, include_cachified=true, use_tuples=false);
     excluded=[:jacobian],  # TODO: figure out why this fails
     logging=LOGGING,
 );
diff --git a/DifferentiationInterface/test/Back/Zygote/test.jl b/DifferentiationInterface/test/Back/Zygote/test.jl
@@ -27,7 +27,9 @@ end
 @testset "Dense" begin
     test_differentiation(
         backends,
-        default_scenarios(; include_constantified=true, include_cachified=true);
+        default_scenarios(;
+            include_constantified=true, include_cachified=true, use_tuples=true
+        );
         excluded=[:second_derivative],
         logging=LOGGING,
     )
diff --git a/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl b/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl
@@ -86,7 +86,9 @@ end
         MyAutoSparse.(
             vcat(adaptive_backends, MixedMode(adaptive_backends[1], adaptive_backends[2]))
         ),
-        sparse_scenarios(; include_constantified=true, include_cachified=true);
+        sparse_scenarios(;
+            include_constantified=true, include_cachified=true, use_tuples=true
+        );
         sparsity=true,
         logging=LOGGING,
     )
diff --git a/DifferentiationInterfaceTest/src/scenarios/default.jl b/DifferentiationInterfaceTest/src/scenarios/default.jl
@@ -559,6 +559,7 @@ function default_scenarios(;
     include_closurified=false,
     include_constantified=false,
     include_cachified=false,
+    use_tuples=false,
 )
     x_ = 0.42
     dx_ = 3.14
@@ -635,7 +636,7 @@ function default_scenarios(;
     include_normal && append!(final_scens, scens)
     include_closurified && append!(final_scens, closurify(scens))
     include_constantified && append!(final_scens, constantify(scens))
-    include_cachified && append!(final_scens, cachify(scens))
+    include_cachified && append!(final_scens, cachify(scens; use_tuples=use_tuples))
 
     return final_scens
 end
diff --git a/DifferentiationInterfaceTest/src/scenarios/modify.jl b/DifferentiationInterfaceTest/src/scenarios/modify.jl
@@ -178,8 +178,12 @@ end
 
 Base.show(io::IO, f::StoreInCache) = print(io, "StoreInCache($(f.f))")
 
-function (sc::StoreInCache{:out})(x, y_cache_tup)
-    y_cache = y_cache_tup.cache
+(sc::StoreInCache{:out})(x, y_cache::NamedTuple) = sc(x, y_cache.useful_cache)
+(sc::StoreInCache{:in})(y, x, y_cache::NamedTuple) = sc(y, x, y_cache.useful_cache)
+(sc::StoreInCache{:out})(x, y_cache::Tuple) = sc(x, first(y_cache))
+(sc::StoreInCache{:in})(y, x, y_cache::Tuple) = sc(y, x, first(y_cache))
+
+function (sc::StoreInCache{:out})(x, y_cache::AbstractArray)
     y = sc.f(x)
     if y isa Number
         y_cache[1] = y
@@ -190,8 +194,7 @@ function (sc::StoreInCache{:out})(x, y_cache_tup)
     end
 end
 
-function (sc::StoreInCache{:in})(y, x, y_cache_tup)
-    y_cache = y_cache_tup.cache
+function (sc::StoreInCache{:in})(y, x, y_cache::AbstractArray)
     sc.f(y_cache, x)
     copyto!(y, y_cache)
     return nothing
@@ -204,14 +207,22 @@ Return a new `Scenario` identical to `scen` except for the function `f`, which i
 
 If `tup=true` the cache is a tuple of arrays, otherwise just an array.
 """
-function cachify(scen::Scenario{op,pl_op,pl_fun}) where {op,pl_op,pl_fun}
+function cachify(scen::Scenario{op,pl_op,pl_fun}; use_tuples) where {op,pl_op,pl_fun}
     (; f,) = scen
     @assert isempty(scen.contexts)
     cache_f = StoreInCache{pl_fun}(f)
-    y_cache = if scen.y isa Number
-        (; cache=[myzero(scen.y)])
+    if use_tuples
+        y_cache = if scen.y isa Number
+            (; useful_cache=[myzero(scen.y)], useless_cache=([myzero(scen.y)],))
+        else
+            (; useful_cache=mysimilar(scen.y), useless_cache=(mysimilar(scen.y),))
+        end
     else
-        (; cache=mysimilar(scen.y))
+        y_cache = if scen.y isa Number
+            [myzero(scen.y)]
+        else
+            mysimilar(scen.y)
+        end
     end
     return Scenario{op,pl_op,pl_fun}(
         cache_f;
@@ -221,7 +232,7 @@ function cachify(scen::Scenario{op,pl_op,pl_fun}) where {op,pl_op,pl_fun}
         contexts=(Cache(y_cache),),
         res1=scen.res1,
         res2=scen.res2,
-        smaller=isnothing(scen.smaller) ? nothing : cachify(scen.smaller),
+        smaller=isnothing(scen.smaller) ? nothing : cachify(scen.smaller; use_tuples),
         name=isnothing(scen.name) ? nothing : scen.name * " [cachified]",
     )
 end
@@ -233,7 +244,7 @@ end
 
 closurify(scens::AbstractVector{<:Scenario}) = closurify.(scens)
 constantify(scens::AbstractVector{<:Scenario}) = constantify.(scens)
-cachify(scens::AbstractVector{<:Scenario}) = cachify.(scens)
+cachify(scens::AbstractVector{<:Scenario}; use_tuples) = cachify.(scens; use_tuples)
 
 function set_smaller(
     scen::Scenario{op,pl_op,pl_fun}, smaller::Scenario
diff --git a/DifferentiationInterfaceTest/src/scenarios/sparse.jl b/DifferentiationInterfaceTest/src/scenarios/sparse.jl
@@ -325,7 +325,10 @@ end
 Create a vector of [`Scenario`](@ref)s with sparse array types, focused on sparse Jacobians and Hessians.
 """
 function sparse_scenarios(;
-    band_sizes=[5, 10, 20], include_constantified=false, include_cachified=false
+    band_sizes=[5, 10, 20],
+    include_constantified=false,
+    include_cachified=false,
+    use_tuples=false,
 )
     x_6 = float.(1:6)
     x_2_3 = float.(reshape(1:6, 2, 3))
@@ -347,6 +350,6 @@ function sparse_scenarios(;
     final_scens = Scenario[]
     append!(final_scens, scens)
     include_constantified && append!(final_scens, constantify(scens))
-    include_cachified && append!(final_scens, cachify(scens))
+    include_cachified && append!(final_scens, cachify(scens; use_tuples))
     return final_scens
 end
diff --git a/DifferentiationInterfaceTest/test/standard.jl b/DifferentiationInterfaceTest/test/standard.jl
@@ -33,7 +33,7 @@ sparse_backend = AutoSparse(
 
 test_differentiation(
     sparse_backend,
-    sparse_scenarios(; include_constantified=true);
+    sparse_scenarios(; include_cachified=true);
     sparsity=true,
     logging=LOGGING,
 )
diff --git a/DifferentiationInterfaceTest/test/weird.jl b/DifferentiationInterfaceTest/test/weird.jl
@@ -28,12 +28,14 @@ gpu_scenarios(;
     include_closurified=true,
     include_batchified=true,
     include_cachified=true,
+    use_tuples=true,
 )
 static_scenarios(;
     include_constantified=true,
     include_closurified=true,
     include_batchified=true,
     include_cachified=true,
+    use_tuples=false,
 )
 
 ## Weird arrays
@@ -54,7 +56,10 @@ test_differentiation(AutoZygote(), gpu_scenarios(); excluded=SECOND_ORDER, loggi
 test_differentiation(
     AutoFiniteDiff(),
     default_scenarios(;
-        include_normal=false, include_closurified=true, include_cachified=true
+        include_normal=false,
+        include_closurified=true,
+        include_cachified=true,
+        use_tuples=true,
     );
     excluded=SECOND_ORDER,
     logging=LOGGING,

Original file line number	Diff line number	Diff line change
`@@ -55,7 +55,7 @@ end;`
`55`	`55`
`56`	`56`	`test_differentiation(`
`57`	`57`	`backends[2],`
`58`		`- default_scenarios(; include_normal=false, include_cachified=true);`
	`58`	`+ default_scenarios(; include_normal=false, include_cachified=true, use_tuples=true);`
`59`	`59`	`excluded=SECOND_ORDER,`
`60`	`60`	`logging=LOGGING,`
`61`	`61`	`)`