Put testing utilities in extension

Project.toml

@@ -17,8 +17,11 @@ DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
@@ -36,6 +39,7 @@ DifferentiationInterfacePolyesterForwardDiffExt = [
     "DiffResults",
 ]
 DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"]
+DifferentiationInterfaceTestExt = ["ForwardDiff", "JET", "Random", "Test"]
 DifferentiationInterfaceZygoteExt = ["Zygote"]
 
 [compat]
@@ -49,8 +53,11 @@ Enzyme = "0.11"
 FillArrays = "1"
 FiniteDiff = "2.22"
 ForwardDiff = "0.10"
+JET = "0.8"
 LinearAlgebra = "1"
 PolyesterForwardDiff = "0.1"
+Random = "1"
 ReverseDiff = "1.15"
+Test = "1"
 Zygote = "0.6"
 julia = "1.10"

docs/src/api.md

@@ -64,6 +64,6 @@ These are not part of the public API.
 
 ```@autodocs
 Modules = [DifferentiationInterface]
-Pages = ["backends.jl", "mode.jl", "utils.jl"]
+Pages = ["backends.jl", "mode.jl", "utils.jl", "DifferentiationTest.jl"]
 Public = false
 ```

ext/DifferentiationInterfaceTestExt/DifferentiationInterfaceTestExt.jl

@@ -0,0 +1,21 @@
+module DifferentiationInterfaceTestExt
+
+# package dependencies
+using ADTypes: AbstractADType
+using DifferentiationInterface
+using DifferentiationInterface.DifferentiationTest
+using DifferentiationInterface: ForwardMode, ReverseMode, autodiff_mode
+import DifferentiationInterface as DI
+import DifferentiationInterface.DifferentiationTest as DT
+using LinearAlgebra: dot
+
+# new dependencies
+using ForwardDiff: ForwardDiff
+using JET: @test_opt
+using Random: AbstractRNG, default_rng, randn!
+using Test: @test, @testset
+
+include("scenarios.jl")
+include("test.jl")
+
+end

ext/DifferentiationInterfaceTestExt/scenarios.jl

@@ -0,0 +1,83 @@
+
+in_type(::Scenario{F,X}) where {F,X} = X
+out_type(::Scenario{F,X,Y}) where {F,X,Y} = Y
+
+function DT.make_scenario(rng::AbstractRNG, f, x)
+    y = f(x)
+    return make_scenario(rng, f, x, y)
+end
+
+## Auto fill
+
+function DT.make_scenario(rng::AbstractRNG, f, x::Number, y::Number)
+    dx = randn(rng, typeof(x))
+    dy = randn(rng, typeof(y))
+    der_true = ForwardDiff.derivative(f, x)
+    dx_true = der_true * dy
+    dy_true = der_true * dx
+    return Scenario(; f, x, y, dx, dy, dx_true, dy_true, der_true)
+end
+
+function DT.make_scenario(rng::AbstractRNG, f, x::Number, y::AbstractArray)
+    dx = randn(rng, typeof(x))
+    dy = similar(y)
+    randn!(rng, dy)
+    multider_true = ForwardDiff.derivative(f, x)
+    dx_true = dot(multider_true, dy)
+    dy_true = multider_true .* dx
+    return Scenario(; f, x, y, dx, dy, dx_true, dy_true, multider_true)
+end
+
+function DT.make_scenario(rng::AbstractRNG, f, x::AbstractArray, y::Number)
+    dx = similar(x)
+    randn!(rng, dx)
+    dy = randn(rng, typeof(y))
+    grad_true = ForwardDiff.gradient(f, x)
+    dx_true = grad_true .* dy
+    dy_true = dot(grad_true, dx)
+    return Scenario(; f, x, y, dx, dy, dx_true, dy_true, grad_true)
+end
+
+function DT.make_scenario(rng::AbstractRNG, f, x::AbstractArray, y::AbstractArray)
+    dx = similar(x)
+    randn!(rng, dx)
+    dy = similar(y)
+    randn!(rng, dy)
+    jac_true = ForwardDiff.jacobian(f, x)
+    dx_true = reshape(transpose(jac_true) * vec(dy), size(x))
+    dy_true = reshape(jac_true * vec(dx), size(y))
+    return Scenario(; f, x, y, dx, dy, dx_true, dy_true, jac_true)
+end
+
+## Defaults
+
+f_scalar_scalar(x::Number)::Number = sin(x)
+
+f_scalar_vector(x::Number)::AbstractVector = [sin(x), sin(2x)]
+f_scalar_matrix(x::Number)::AbstractMatrix = hcat([sin(x) cos(x)], [sin(2x) cos(2x)])
+
+f_vector_scalar(x::AbstractVector)::Number = sum(sin, x)
+f_matrix_scalar(x::AbstractMatrix)::Number = sum(sin, x)
+
+f_vector_vector(x::AbstractVector)::AbstractVector = vcat(sin.(x), cos.(x))
+f_vector_matrix(x::AbstractVector)::AbstractMatrix = hcat(sin.(x), cos.(x))
+
+f_matrix_vector(x::AbstractMatrix)::AbstractVector = vcat(vec(sin.(x)), vec(cos.(x)))
+f_matrix_matrix(x::AbstractMatrix)::AbstractMatrix = hcat(vec(sin.(x)), vec(cos.(x)))
+
+function DT.default_scenarios(rng::AbstractRNG)
+    scenarios = [
+        make_scenario(rng, f_scalar_scalar, 1.0),
+        make_scenario(rng, f_scalar_vector, 1.0),
+        make_scenario(rng, f_scalar_matrix, 1.0),
+        make_scenario(rng, f_vector_scalar, [1.0, 2.0]),
+        make_scenario(rng, f_matrix_scalar, [1.0 2.0; 3.0 4.0]),
+        make_scenario(rng, f_vector_vector, [1.0, 2.0]),
+        make_scenario(rng, f_vector_matrix, [1.0, 2.0]),
+        make_scenario(rng, f_matrix_vector, [1.0 2.0; 3.0 4.0]),
+        make_scenario(rng, f_matrix_matrix, [1.0 2.0; 3.0 4.0]),
+    ]
+    return scenarios
+end
+
+DT.default_scenarios() = default_scenarios(default_rng())