Replace AD testing utilities with DifferentiationInterfaceTest.jl

test/Project.toml

@@ -1,40 +1,42 @@
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AxisArrays = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
+DifferentiationInterfaceTest = "a82114a7-5aa3-49a8-9643-716bb13727a3"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 Kronecker = "2c470bb0-bcc8-11e8-3dad-c9649493f05e"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
-StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
+ADTypes = "1"
 AxisArrays = "0.4.3"
 Compat = "3, 4"
+DifferentiationInterface = "0.7"
+DifferentiationInterfaceTest = "0.11"
 Distances = "0.10"
 Documenter = "0.25, 0.26, 0.27, 1"
 FiniteDifferences = "0.10.8, 0.11, 0.12"
 ForwardDiff = "0.10"
-Functors = "0.2, 0.3, 0.4, 0.5"
 Kronecker = "0.4, 0.5"
 LogExpFunctions = "0.2, 0.3"
 PDMats = "0.9, 0.10, 0.11"
 ReverseDiff = "1.2"
 SpecialFunctions = "0.10, 1, 2"
-StableRNGs = "1"
 StaticArrays = "1"
 Statistics = "1"
 Zygote = "0.6.38"

test/basekernels/constant.jl

@@ -10,7 +10,6 @@
         TestUtils.test_interface(k, Float64)
         TestUtils.test_interface(k, Vector{String})
         test_ADs(ZeroKernel)
-        test_interface_ad_perf(_ -> k, nothing, StableRNG(123456))
     end
     @testset "WhiteKernel" begin
         k = WhiteKernel()
@@ -25,7 +24,6 @@
         TestUtils.test_interface(k, Float64)
         TestUtils.test_interface(k, Vector{String})
         test_ADs(WhiteKernel)
-        test_interface_ad_perf(_ -> k, nothing, StableRNG(123456))
     end
     @testset "ConstantKernel" begin
         c = 2.0
@@ -36,12 +34,9 @@
         @test metric(ConstantKernel()) == KernelFunctions.Delta()
         @test metric(ConstantKernel(; c=2.0)) == KernelFunctions.Delta()
         @test repr(k) == "Constant Kernel (c = $(c))"
-        test_params(k, ([c],))
-
         # Standardised tests.
         TestUtils.test_interface(k, Float64)
         TestUtils.test_interface(k, Vector{String})
         test_ADs(c -> ConstantKernel(; c=only(c)), [c])
-        test_interface_ad_perf(c -> ConstantKernel(; c=c), c, StableRNG(123456))
     end
 end

test/basekernels/cosine.jl

@@ -20,5 +20,4 @@
     # Standardised tests.
     TestUtils.test_interface(k, Vector{Float64})
     test_ADs(CosineKernel)
-    test_interface_ad_perf(_ -> CosineKernel(), nothing, StableRNG(123456))
 end

test/basekernels/exponential.jl

@@ -22,7 +22,6 @@
         # Standardised tests.
         TestUtils.test_interface(k)
         test_ADs(SEKernel)
-        test_interface_ad_perf(_ -> SEKernel(), nothing, StableRNG(123456))
     end
     @testset "ExponentialKernel" begin
         k = ExponentialKernel()
@@ -41,7 +40,6 @@
         # Standardised tests.
         TestUtils.test_interface(k)
         test_ADs(ExponentialKernel)
-        test_interface_ad_perf(_ -> ExponentialKernel(), nothing, StableRNG(123456))
     end
     @testset "GammaExponentialKernel" begin
         γ = 1.0
@@ -59,7 +57,6 @@
         @test k2(v1, v2) ≈ k(v1, v2)
 
         test_ADs(γ -> GammaExponentialKernel(; gamma=only(γ)), [1 + 0.5 * rand()])
-        test_params(k, ([γ],))
         TestUtils.test_interface(GammaExponentialKernel(; γ=1.36))
 
         #Coherence :

test/basekernels/exponentiated.jl

@@ -14,5 +14,4 @@
     # Standardised tests. This kernel appears to be fairly numerically unstable.
     TestUtils.test_interface(k; atol=1e-3)
     test_ADs(ExponentiatedKernel)
-    test_interface_ad_perf(_ -> ExponentiatedKernel(), nothing, StableRNG(123456))
 end

test/basekernels/fbm.jl

@@ -23,8 +23,4 @@
             Zygote.gradient((x, y) -> sum(f.(x, y)), zeros(1), fill(0.9, 1))[1][1]
         )
     end
-
-    test_params(k, ([h],))
-
-    test_interface_ad_perf(h -> FBMKernel(; h=h), h, StableRNG(123456))
 end

test/basekernels/gabor.jl

@@ -32,10 +32,4 @@
         ),
         [ell, p],
     )
-    test_interface_ad_perf((ell, p), StableRNG(123456)) do θ
-        gaborkernel(;
-            sqexponential_transform=ScaleTransform(θ[1]),
-            cosine_transform=ScaleTransform(θ[2]),
-        )
-    end
 end

test/basekernels/matern.jl

@@ -21,29 +21,6 @@
         # Standardised tests.
         TestUtils.test_interface(k, Float64)
         test_ADs(() -> MaternKernel(; nu=ν))
-
-        test_params(k, ([ν],))
-
-        # The performance of this kernel varies quite a lot from method to method, so
-        # requires us to specify whether performance tests pass or not.
-        @testset "performance ($T)" for T in [
-            Vector{Float64},
-            ColVecs{Float64,Matrix{Float64}},
-            RowVecs{Float64,Matrix{Float64}},
-        ]
-            xs = example_inputs(StableRNG(123456), Vector{Float64})
-            test_interface_ad_perf(
-                ν -> MaternKernel(; nu=ν),
-                ν,
-                xs...;
-                passes=(
-                    unary=(false, false, false),
-                    binary=(false, false, false),
-                    diag_unary=(true, false, false),
-                    diag_binary=(true, false, false),
-                ),
-            )
-        end
     end
     @testset "Matern32Kernel" begin
         k = Matern32Kernel()
@@ -60,7 +37,6 @@
         # Standardised tests.
         TestUtils.test_interface(k, Float64)
         test_ADs(Matern32Kernel)
-        test_interface_ad_perf(_ -> Matern32Kernel(), nothing, StableRNG(123456))
     end
     @testset "Matern52Kernel" begin
         k = Matern52Kernel()
@@ -80,7 +56,6 @@
         # Standardised tests.
         TestUtils.test_interface(k, Float64)
         test_ADs(Matern52Kernel)
-        test_interface_ad_perf(_ -> Matern52Kernel(), nothing, StableRNG(123456))
     end
     @testset "Matern72Kernel" begin
         k = Matern72Kernel()
@@ -104,7 +79,6 @@
         # Standardised tests.
         TestUtils.test_interface(k, Float64)
         test_ADs(Matern72Kernel)
-        test_interface_ad_perf(_ -> Matern72Kernel(), nothing, StableRNG(123456))
     end
     @testset "Coherence Materns" begin
         @test kappa(MaternKernel(; ν=0.5), x) ≈ kappa(ExponentialKernel(), x)

test/basekernels/nn.jl

@@ -8,5 +8,4 @@
     # Standardised tests.
     TestUtils.test_interface(k, Float64)
     test_ADs(NeuralNetworkKernel)
-    test_interface_ad_perf(_ -> NeuralNetworkKernel(), nothing, StableRNG(123456))
 end

test/basekernels/periodic.jl

@@ -17,5 +17,4 @@
     TestUtils.test_interface(PeriodicKernel(; r=[0.8, 0.7]), RowVecs{Float64})
 
     test_ADs(r -> PeriodicKernel(; r=exp.(r)), log.(r))
-    test_params(k, (r,))
 end

test/basekernels/piecewisepolynomial.jl

@@ -32,29 +32,4 @@
     TestUtils.test_interface(k, ColVecs{Float64}; dim_in=2)
     TestUtils.test_interface(k, RowVecs{Float64}; dim_in=2)
     test_ADs(() -> PiecewisePolynomialKernel{degree}(; dim=D))
-
-    test_params(k, ())
-    if VERSION >= v"1.8.0"
-        test_interface_ad_perf(nothing, StableRNG(123456)) do _
-            PiecewisePolynomialKernel{degree}(; dim=D)
-        end
-    else
-        @testset "AD Alloc Performance ($T)" for T in [
-            Vector{Float64},
-            ColVecs{Float64,Matrix{Float64}},
-            RowVecs{Float64,Matrix{Float64}},
-        ]
-            test_interface_ad_perf(
-                _ -> PiecewisePolynomialKernel{degree}(; dim=D),
-                nothing,
-                example_inputs(StableRNG(123456), T)...;
-                passes=(
-                    unary=(true, true, true),
-                    binary=(true, true, true),
-                    diag_unary=(true, true, true),
-                    diag_binary=(true, true, true),
-                ),
-            )
-        end
-    end
 end

test/basekernels/polynomial.jl

@@ -19,8 +19,6 @@
         # Standardised tests.
         TestUtils.test_interface(k, Float64)
         test_ADs(x -> LinearKernel(; c=x[1]), [c])
-        test_params(LinearKernel(; c=c), ([c],))
-        test_interface_ad_perf(c -> LinearKernel(; c=c), c, StableRNG(123456))
     end
     @testset "PolynomialKernel" begin
         k = PolynomialKernel()
@@ -42,9 +40,5 @@
         # Standardised tests.
         TestUtils.test_interface(k, Float64)
         test_ADs(x -> PolynomialKernel(; c=x[1]), [c])
-        test_params(PolynomialKernel(; c=c), ([c],))
-        test_interface_ad_perf(
-            c -> PolynomialKernel(; degree=2, c=c), 0.3, StableRNG(123456)
-        )
     end
 end

test/basekernels/rational.jl

@@ -28,8 +28,6 @@
         # Standardised tests.
         TestUtils.test_interface(k, Float64)
         test_ADs(x -> RationalKernel(; alpha=exp(x[1])), [α])
-        test_params(k, ([α],))
-        test_interface_ad_perf(α -> RationalKernel(; alpha=α), α, StableRNG(123456))
     end
 
     @testset "RationalQuadraticKernel" begin
@@ -56,12 +54,8 @@
         # Standardised tests.
         TestUtils.test_interface(k, Float64)
         # test_ADs(x -> RationalQuadraticKernel(; alpha=exp(x[1])), [α])
-        test_params(k, ([α],))
-        test_interface_ad_perf(α, StableRNG(123456)) do α
-            RationalQuadraticKernel(; alpha=α)
-        end
 
-        # Check correctness and performance with non-Euclidean metrics.
+        # Check correctness with non-Euclidean metrics.
         TestUtils.test_interface(
             RationalQuadraticKernel(; alpha=α, metric=WeightedEuclidean([1.0, 2.0])),
             ColVecs{Float64},
@@ -70,10 +64,6 @@
             RationalQuadraticKernel(; alpha=α, metric=WeightedEuclidean([1.0, 2.0])),
             RowVecs{Float64},
         )
-        types = [ColVecs{Float64,Matrix{Float64}}, RowVecs{Float64,Matrix{Float64}}]
-        test_interface_ad_perf(α, StableRNG(123456)) do α
-            RationalQuadraticKernel(; alpha=α, metric=KernelFunctions.DotProduct())
-        end
     end
 
     @testset "GammaRationalKernel" begin
@@ -147,9 +137,5 @@
         TestUtils.test_interface(k, Float64)
         a = 1.0 + rand()
         test_ADs(x -> GammaRationalKernel(; α=x[1], γ=x[2]), [a, 1 + 0.5 * rand()])
-        test_params(GammaRationalKernel(; α=a, γ=x), ([a], [x]))
-        test_interface_ad_perf((2.0, 1.5), StableRNG(123456)) do θ
-            GammaRationalKernel(; α=θ[1], γ=θ[2])
-        end
     end
 end

test/kernels/kernelproduct.jl

@@ -21,11 +21,6 @@
     test_ADs(
         x -> KernelProduct(SqExponentialKernel(), LinearKernel(; c=exp(x[1]))), rand(1)
     )
-    test_interface_ad_perf(2.4, StableRNG(123456)) do c
-        KernelProduct(SqExponentialKernel(), LinearKernel(; c=c))
-    end
-    test_params(k1 * k2, (k1, k2))
-
     nested_k = RBFKernel() * (LinearKernel() + CosineKernel() * RBFKernel())
     test_type_stability(nested_k)
 end

test/kernels/kernelsum.jl

@@ -22,11 +22,6 @@
     test_interface(kvec, Float64)
     test_interface(ConstantKernel(; c=1.5) + WhiteKernel(), Vector{String})
     test_ADs(x -> KernelSum(SqExponentialKernel(), LinearKernel(; c=exp(x[1]))), rand(1))
-    test_interface_ad_perf(2.4, StableRNG(123456)) do c
-        KernelSum(SqExponentialKernel(), LinearKernel(; c=c))
-    end
-
-    test_params(k1 + k2, (k1, k2))
 
     # Regression tests for https://github.com//issues/458
     @testset for k in (

test/kernels/kerneltensorproduct.jl

@@ -46,12 +46,6 @@
         rand(1);
         dims=[2, 2],
     )
-    types = [ColVecs{Float64,Matrix{Float64}}, RowVecs{Float64,Matrix{Float64}}]
-    test_interface_ad_perf(2.1, StableRNG(123456), types) do c
-        KernelTensorProduct(SqExponentialKernel(), LinearKernel(; c=c))
-    end
-    test_params(KernelTensorProduct(k1, k2), (k1, k2))
-
     @testset "single kernel" begin
         kernel = KernelTensorProduct(k1)
         @test length(kernel) == 1

test/kernels/normalizedkernel.jl

@@ -11,9 +11,4 @@
     # Standardised tests.
     TestUtils.test_interface(kn, Float64)
     test_ADs(x -> NormalizedKernel(exp(x[1]) * SqExponentialKernel()), rand(1))
-    test_interface_ad_perf(0.3, StableRNG(123456)) do c
-        NormalizedKernel(c * SqExponentialKernel())
-    end
-
-    test_params(kn, k)
 end

test/kernels/scaledkernel.jl

@@ -12,7 +12,4 @@
     # Standardised tests.
     TestUtils.test_interface(ks, Float64)
     test_ADs(x -> exp(x[1]) * SqExponentialKernel(), rand(1))
-    test_interface_ad_perf(c -> c * SEKernel(), 0.3, StableRNG(123456))
-
-    test_params(s * k, (k, [s]))
 end

test/kernels/transformedkernel.jl

@@ -30,46 +30,4 @@
         Vector{String},
     )
     test_ADs(x -> SqExponentialKernel() ∘ ScaleTransform(x[1]), rand(1))
-    test_interface_ad_perf(0.35, StableRNG(123456)) do λ
-        SqExponentialKernel() ∘ ScaleTransform(λ)
-    end
-
-    # Test implicit gradients
-    @testset "Implicit gradients" begin
-        if _TEST_ZYGOTE
-            k = SqExponentialKernel() ∘ ScaleTransform(2.0)
-            ps = params(k)
-            X = rand(10, 1)
-            x = vec(X)
-            A = rand(10, 10)
-            # Implicit
-            g1 = Zygote.gradient(ps) do
-                tr(kernelmatrix(k, X; obsdim=1) * A)
-            end
-            # Explicit
-            g2 = Zygote.gradient(k) do k
-                tr(kernelmatrix(k, X; obsdim=1) * A)
-            end
-
-            # Implicit for a vector
-            g3 = Zygote.gradient(ps) do
-                tr(kernelmatrix(k, x) * A)
-            end
-            @test g1[first(ps)] ≈ first(g2).transform.s
-            @test g1[first(ps)] ≈ g3[first(ps)]
-        else
-            @test_broken false  # Zygote not supported on Julia >= 1.12
-        end
-    end
-
-    @testset "Parameters" begin
-        k = ConstantKernel(; c=rand(rng))
-        # c = Chain(Dense(3, 2))
-
-        test_params(k ∘ ScaleTransform(s), (k, [s]))
-        test_params(k ∘ ARDTransform(v), (k, v))
-        test_params(k ∘ LinearTransform(P), (k, P))
-        test_params(k ∘ LinearTransform(P) ∘ ScaleTransform(s), (k, [s], P))
-        # test_params(k ∘ FunctionTransform(c), (k, c))
-    end
 end