CompatHelper: bump compat for "ChainRulesCore" to "1" (#344)

* CompatHelper: bump compat for "ChainRulesCore" to "1"

* bump patch version

* Update InplaceableThunks

* Restrict to version 1

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Remove Flux dependency

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Remove another Flux reference

* FD + Zygote issue

* Tweak fbm tests

* Note Flux code

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: ST John <st--@users.noreply.github.com>
Co-authored-by: WT <wt0881@my.bristol.ac.uk>
Co-authored-by: willtebbutt <wct23@cam.ac.uk>

Author: 3 people

Project.toml

@@ -1,6 +1,6 @@
 name = "KernelFunctions"
 uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
-version = "0.10.13"
+version = "0.10.14"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -21,7 +21,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
-ChainRulesCore = "0.9.44, 0.10"
+ChainRulesCore = "1"
 Compat = "3.7"
 CompositionsBase = "0.1"
 Distances = "0.10"

src/chainrules.jl

@@ -127,13 +127,13 @@ function ChainRulesCore.rrule(
     function SqMahalanobis_pullback(Δ::Real)
         a_b = a - b
         ∂qmat = InplaceableThunk(
-            @thunk((a_b * a_b') * Δ), X̄ -> mul!(X̄, a_b, a_b', true, Δ)
+            X̄ -> mul!(X̄, a_b, a_b', true, Δ), @thunk((a_b * a_b') * Δ)
         )
         ∂a = InplaceableThunk(
-            @thunk((2 * Δ) * dist.qmat * a_b), X̄ -> mul!(X̄, dist.qmat, a_b, true, 2 * Δ)
+            X̄ -> mul!(X̄, dist.qmat, a_b, true, 2 * Δ), @thunk((2 * Δ) * dist.qmat * a_b)
         )
         ∂b = InplaceableThunk(
-            @thunk((-2 * Δ) * dist.qmat * a_b), X̄ -> mul!(X̄, dist.qmat, a_b, true, -2 * Δ)
+            X̄ -> mul!(X̄, dist.qmat, a_b, true, -2 * Δ), @thunk((-2 * Δ) * dist.qmat * a_b)
         )
         return Tangent{typeof(dist)}(; qmat=∂qmat), ∂a, ∂b
     end

test/Project.toml

@@ -3,8 +3,8 @@ AxisArrays = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
-Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 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"
@@ -20,8 +20,8 @@ AxisArrays = "0.4.3"
 Distances = "0.9, 0.10"
 Documenter = "0.25, 0.26, 0.27"
 FiniteDifferences = "0.10.8, 0.11, 0.12"
-Flux = "0.10, 0.11, 0.12"
 ForwardDiff = "0.10"
+Functors = "0.2"
 Kronecker = "0.4"
 LogExpFunctions = "0.2, 0.3"
 PDMats = "0.9, 0.10, 0.11"

test/basekernels/fbm.jl

@@ -15,17 +15,12 @@
     @test repr(k) == "Fractional Brownian Motion Kernel (h = $(h))"
     test_ADs(FBMKernel; ADs=[:ReverseDiff])
 
-    # Tests failing for ForwardDiff and Zygote@0.6 (obtained with Julia > 1.3).
+    # Tests failing for ForwardDiff and Zygote@0.6.
     # Related to: https://github.com/FluxML/Zygote.jl/issues/1036
-    @test_broken !isinf(ForwardDiff.gradient(x -> x[1]^x[2], [0.0, 0.9])[1])
-    if VERSION >= v"1.4.0"
-        f(x, y) = x^y
-        @test_broken !isinf(
-            Zygote.gradient((x, y) -> sum(f.(x, y)), zeros(1), fill(0.9, 1))[1][1]
-        )
-    else
-        test_ADs(FBMKernel; ADs=[:Zygote])
-    end
+    f(x, y) = x^y
+    @test_broken !isinf(
+        Zygote.gradient((x, y) -> sum(f.(x, y)), zeros(1), fill(0.9, 1))[1][1]
+    )
 
     test_params(k, ([h],))
 end

test/kernels/neuralkernelnetwork.jl

@@ -19,7 +19,7 @@ using KernelFunctions: NeuralKernelNetwork, LinearLayer, product, Primitive
         primitives = Primitive(k1, k2)
 
         # Build NKN Kernel.
-        nkn = NeuralKernelNetwork(primitives, Chain(LinearLayer(2, 2), product))
+        nkn = NeuralKernelNetwork(primitives, product)
 
         # Apply standard test suite.
         TestUtils.test_interface(nkn, Float64)

test/kernels/transformedkernel.jl

@@ -30,21 +30,21 @@
     # Test implicit gradients
     @testset "Implicit gradients" begin
         k = SqExponentialKernel() ∘ ScaleTransform(2.0)
-        ps = Flux.params(k)
+        ps = params(k)
         X = rand(10, 1)
         x = vec(X)
         A = rand(10, 10)
         # Implicit
-        g1 = Flux.gradient(ps) do
+        g1 = Zygote.gradient(ps) do
             tr(kernelmatrix(k, X; obsdim=1) * A)
         end
         # Explicit
-        g2 = Flux.gradient(k) do k
+        g2 = Zygote.gradient(k) do k
             tr(kernelmatrix(k, X; obsdim=1) * A)
         end
 
         # Implicit for a vector
-        g3 = Flux.gradient(ps) do
+        g3 = Zygote.gradient(ps) do
             tr(kernelmatrix(k, x) * A)
         end
         @test g1[first(ps)] ≈ first(g2).transform.s
@@ -53,12 +53,12 @@
 
     @testset "Parameters" begin
         k = ConstantKernel(; c=rand(rng))
-        c = Chain(Dense(3, 2))
+        # 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))
+        # test_params(k ∘ FunctionTransform(c), (k, c))
     end
 end

test/runtests.jl

@@ -2,14 +2,14 @@ using KernelFunctions
 using AxisArrays
 using Distances
 using Documenter
+using Functors: functor
 using Kronecker: Kronecker
 using LinearAlgebra
 using LogExpFunctions
 using PDMats
 using Random
 using SpecialFunctions
 using Test
-using Flux
 using Zygote: Zygote
 using ForwardDiff: ForwardDiff
 using ReverseDiff: ReverseDiff

test/test_utils.jl

@@ -1,7 +1,25 @@
 # More test utilities. Can't be included in KernelFunctions because they introduce a number
 # of additional deps that we don't want to have in the main package.
 
-# Check parameters of kernels
+# Check parameters of kernels. `trainable`, `params!`, and `params` are taken directly from
+# Flux.jl so as to avoid having to depend on Flux at test-time.
+trainable(m) = functor(m)[1]
+
+params!(p::Zygote.Params, x::AbstractArray{<:Number}, seen=Zygote.IdSet()) = push!(p, x)
+
+function params!(p::Zygote.Params, x, seen=Zygote.IdSet())
+    x in seen && return nothing
+    push!(seen, x)
+    for child in trainable(x)
+        params!(p, child, seen)
+    end
+end
+
+function params(m...)
+    ps = Zygote.Params()
+    params!(ps, m)
+    return ps
+end
 
 function test_params(kernel, reference)
     params_kernel = params(kernel)