fix: refactor test loops (#848)

* Improve type stability in correctness tests

* Add some fixes

* Fixes

* Missing conj

* Reprepare cov

* Add changelog

Author: gdalle

DifferentiationInterfaceTest/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterfaceTest-v0.10.1...main)
 
+### Fixed
+
+  - Refactor test loops ([#848](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/848))
+
 ## [0.10.1](https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterfaceTest-v0.10.0...DifferentiationInterfaceTest-v0.10.1)
 
 ### Added

DifferentiationInterfaceTest/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterfaceTest"
 uuid = "a82114a7-5aa3-49a8-9643-716bb13727a3"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.10.1"
+version = "0.10.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -12,6 +12,7 @@ DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -54,6 +55,7 @@ JLArrays = "0.1, 0.2"
 LinearAlgebra = "1"
 Lux = "1.1.0"
 LuxTestUtils = "1.3.1, 2"
+PrecompileTools = "1.2.1"
 ProgressMeter = "1"
 Random = "1"
 SparseArrays = "1"

DifferentiationInterfaceTest/ext/DifferentiationInterfaceTestJLArraysExt/DifferentiationInterfaceTestJLArraysExt.jl

@@ -3,6 +3,7 @@ module DifferentiationInterfaceTestJLArraysExt
 import DifferentiationInterface as DI
 import DifferentiationInterfaceTest as DIT
 using JLArrays: JLArray, JLVector, JLMatrix, jl
+using PrecompileTools: @compile_workload
 
 jl_num_to_vec(x::Number) = sin.(jl([1, 2]) .* x)
 jl_num_to_mat(x::Number) = hcat(jl_num_to_vec(x), jl_num_to_vec(3x))
@@ -42,4 +43,8 @@ function DIT.gpu_scenarios(args...; kwargs...)
     return myjl.(scens)
 end
 
+@compile_workload begin
+    DIT.gpu_scenarios(; include_constantified=true, include_cachified=true)
+end
+
 end

DifferentiationInterfaceTest/ext/DifferentiationInterfaceTestStaticArraysExt/DifferentiationInterfaceTestStaticArraysExt.jl

@@ -4,6 +4,7 @@ import DifferentiationInterface as DI
 import DifferentiationInterfaceTest as DIT
 using SparseArrays: SparseArrays, SparseMatrixCSC, nnz, spdiagm
 using StaticArrays: StaticArray, MArray, MMatrix, MVector, SArray, SMatrix, SVector
+using PrecompileTools: @compile_workload
 
 static_num_to_vec(x::Number) = sin.(SVector(1, 2) .* x)
 static_num_to_mat(x::Number) = hcat(static_num_to_vec(x), static_num_to_vec(3x))
@@ -61,4 +62,8 @@ function DIT.static_scenarios(args...; kwargs...)
     return mystatic.(scens)
 end
 
+@compile_workload begin
+    DIT.static_scenarios(; include_constantified=true, include_cachified=true)
+end
+
 end

DifferentiationInterfaceTest/src/DifferentiationInterfaceTest.jl

@@ -94,6 +94,7 @@ using DifferentiationInterface: PreparationMismatchError
 using DocStringExtensions: TYPEDFIELDS, TYPEDSIGNATURES
 using JET: @test_opt
 using LinearAlgebra: Adjoint, Diagonal, Transpose, I, dot, parent
+using PrecompileTools: @compile_workload
 using ProgressMeter: ProgressUnknown, next!
 using Random: AbstractRNG, default_rng, rand!
 using SparseArrays:
@@ -114,7 +115,16 @@ List of all second-order operators, to facilitate exclusion during tests.
 """
 const SECOND_ORDER = [:hvp, :second_derivative, :hessian]
 
-const ALL_OPS = vcat(FIRST_ORDER, SECOND_ORDER)
+const ALL_OPS = (
+    :pushforward,
+    :pullback,
+    :derivative,
+    :gradient,
+    :jacobian,
+    :hvp,
+    :second_derivative,
+    :hessian,
+)
 
 include("utils.jl")
 
@@ -128,6 +138,7 @@ include("scenarios/empty.jl")
 include("scenarios/extensions.jl")
 
 include("tests/correctness_eval.jl")
+include("tests/prep_eval.jl")
 include("tests/type_stability_eval.jl")
 include("tests/benchmark.jl")
 include("tests/benchmark_eval.jl")
@@ -140,4 +151,8 @@ export Scenario, compute_results
 export test_differentiation, benchmark_differentiation
 export DifferentiationBenchmarkDataRow
 
+@compile_workload begin
+    default_scenarios(; include_constantified=true, include_cachified=true)
+end
+
 end

DifferentiationInterfaceTest/src/scenarios/default.jl

@@ -306,7 +306,7 @@ function arr_to_num_hessian(x0)
     return convert(typeof(similar(x0, length(x0), length(x0))), H)
 end
 
-arr_to_num_pushforward(x, dx) = sum(arr_to_num_gradient(x) .* dx)
+arr_to_num_pushforward(x, dx) = sum(conj.(arr_to_num_gradient(x)) .* dx)
 arr_to_num_pullback(x, dy) = arr_to_num_gradient(x) .* dy
 arr_to_num_hvp(x, dx) = reshape(arr_to_num_hessian(x) * vec(dx), size(x))
 

DifferentiationInterfaceTest/src/test_differentiation.jl

@@ -164,6 +164,7 @@ function test_differentiation(
                             sparsity,
                             reprepare,
                         )
+                        test_prep(adapted_backend, scen)
                     end
                     yield()
                     (type_stability != :none) && @testset "Type stability" begin