Merge branch 'master' into dependabot/github_actions/actions/upload-artifact-4

Author: simsurace

.github/workflows/benchmark-comment.yml

@@ -41,7 +41,7 @@ jobs:
           echo ::set-output name=body::$(cat ./pull-request-number.artifact)
       # check if the previous comment exists
       - name: find comment
-        uses: peter-evans/find-comment@v1
+        uses: peter-evans/find-comment@v3
         id: fc
         with:
           issue-number: ${{ steps.output-pull-request-number.outputs.body }}

Project.toml

@@ -1,6 +1,6 @@
 name = "KernelFunctions"
 uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
-version = "0.10.60"
+version = "0.10.63"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

README.md

@@ -1,6 +1,6 @@
 # KernelFunctions.jl
 
-![CI](https://github.com/JuliaGaussianProcesses/KernelFunctions.jl/workflows/CI/badge.svg?branch=master)
+[![CI](https://github.com/JuliaGaussianProcesses/KernelFunctions.jl/actions/workflows/ci.yml/badge.svg?branch=master)](https://github.com/JuliaGaussianProcesses/KernelFunctions.jl/actions/workflows/ci.yml?query=branch%3Amaster)
 [![codecov](https://codecov.io/gh/JuliaGaussianProcesses/KernelFunctions.jl/branch/master/graph/badge.svg?token=rmDh3gb7hN)](https://codecov.io/gh/JuliaGaussianProcesses/KernelFunctions.jl)
 [![Documentation (stable)](https://img.shields.io/badge/docs-stable-blue.svg)](https://juliagaussianprocesses.github.io/KernelFunctions.jl/stable)
 [![Documentation (latest)](https://img.shields.io/badge/docs-dev-blue.svg)](https://juliagaussianprocesses.github.io/KernelFunctions.jl/dev)

src/KernelFunctions.jl

@@ -48,7 +48,7 @@ export tensor, ⊗, compose
 
 using Compat
 using ChainRulesCore: ChainRulesCore, Tangent, ZeroTangent, NoTangent
-using ChainRulesCore: @thunk, InplaceableThunk
+using ChainRulesCore: @thunk, InplaceableThunk, ProjectTo, unthunk
 using CompositionsBase
 using Distances
 using FillArrays

src/chainrules.jl

@@ -111,16 +111,113 @@ end
 
 function ChainRulesCore.rrule(s::Sinus, x::AbstractVector, y::AbstractVector)
     d = x - y
-    sind = sinpi.(d)
-    abs2_sind_r = abs2.(sind) ./ s.r
+    abs2_sind_r = (sinpi.(d) ./ s.r) .^ 2
     val = sum(abs2_sind_r)
-    gradx = twoπ .* cospi.(d) .* sind ./ (s.r .^ 2)
+    gradx = π .* sinpi.(2 .* d) ./ s.r .^ 2
     function evaluate_pullback(Δ::Any)
-        return (r=-2Δ .* abs2_sind_r,), Δ * gradx, -Δ * gradx
+        r̄ = -2Δ .* abs2_sind_r ./ s.r
+        s̄ = ChainRulesCore.Tangent{typeof(s)}(; r=r̄)
+        return s̄, Δ * gradx, -Δ * gradx
     end
     return val, evaluate_pullback
 end
 
+function ChainRulesCore.rrule(
+    ::typeof(Distances.pairwise), d::Sinus, x::AbstractMatrix; dims=2
+)
+    project_x = ProjectTo(x)
+    function pairwise_pullback(z̄)
+        Δ = unthunk(z̄)
+        n = size(x, dims)
+        x̄ = collect(zero(x))
+        r̄ = zero(d.r)
+        if dims == 1
+            for j in 1:n, i in 1:n
+                xi = view(x, i, :)
+                xj = view(x, j, :)
+                ds = π .* Δ[i, j] .* sinpi.(2 .* (xi .- xj)) ./ d.r .^ 2
+                r̄ .-= 2 .* Δ[i, j] .* sinpi.(xi .- xj) .^ 2 ./ d.r .^ 3
+                x̄[i, :] += ds
+                x̄[j, :] -= ds
+            end
+        elseif dims == 2
+            for j in 1:n, i in 1:n
+                xi = view(x, :, i)
+                xj = view(x, :, j)
+                ds = twoπ .* Δ[i, j] .* sinpi.(xi .- xj) .* cospi.(xi .- xj) ./ d.r .^ 2
+                r̄ .-= 2 .* Δ[i, j] .* sinpi.(xi .- xj) .^ 2 ./ d.r .^ 3
+                x̄[:, i] .+= ds
+                x̄[:, j] .-= ds
+            end
+        end
+        d̄ = ChainRulesCore.Tangent{typeof(d)}(; r=r̄)
+        return NoTangent(), d̄, @thunk(project_x(x̄))
+    end
+    return Distances.pairwise(d, x; dims), pairwise_pullback
+end
+
+function ChainRulesCore.rrule(
+    ::typeof(Distances.pairwise), d::Sinus, x::AbstractMatrix, y::AbstractMatrix; dims=2
+)
+    project_x = ProjectTo(x)
+    project_y = ProjectTo(y)
+    function pairwise_pullback(z̄)
+        Δ = unthunk(z̄)
+        n = size(x, dims)
+        m = size(y, dims)
+        x̄ = collect(zero(x))
+        ȳ = collect(zero(y))
+        r̄ = zero(d.r)
+        if dims == 1
+            for j in 1:m, i in 1:n
+                xi = view(x, i, :)
+                yj = view(y, j, :)
+                ds = π .* Δ[i, j] .* sinpi.(2 .* (xi .- yj)) ./ d.r .^ 2
+                r̄ .-= 2 .* Δ[i, j] .* sinpi.(xi .- yj) .^ 2 ./ d.r .^ 3
+                x̄[i, :] .+= ds
+                ȳ[j, :] .-= ds
+            end
+        elseif dims == 2
+            for j in 1:m, i in 1:n
+                xi = view(x, :, i)
+                yj = view(y, :, j)
+                ds = π .* Δ[i, j] .* sinpi.(2 .* (xi .- yj)) ./ d.r .^ 2
+                r̄ .-= 2 .* Δ[i, j] .* sinpi.(xi .- yj) .^ 2 ./ d.r .^ 3
+                x̄[:, i] .+= ds
+                ȳ[:, j] .-= ds
+            end
+        end
+        d̄ = ChainRulesCore.Tangent{typeof(d)}(; r=r̄)
+        return NoTangent(), d̄, @thunk(project_x(x̄)), @thunk(project_y(ȳ))
+    end
+    return Distances.pairwise(d, x, y; dims), pairwise_pullback
+end
+
+function ChainRulesCore.rrule(
+    ::typeof(Distances.colwise), d::Sinus, x::AbstractMatrix, y::AbstractMatrix
+)
+    project_x = ProjectTo(x)
+    project_y = ProjectTo(y)
+    function colwise_pullback(z̄)
+        Δ = unthunk(z̄)
+        n = size(x, 2)
+        x̄ = collect(zero(x))
+        ȳ = collect(zero(y))
+        r̄ = zero(d.r)
+        for i in 1:n
+            xi = view(x, :, i)
+            yi = view(y, :, i)
+            ds = π .* Δ[i] .* sinpi.(2 .* (xi .- yi)) ./ d.r .^ 2
+            r̄ .-= 2 .* Δ[i] .* sinpi.(xi .- yi) .^ 2 ./ d.r .^ 3
+            x̄[:, i] .+= ds
+            ȳ[:, i] .-= ds
+        end
+        d̄ = ChainRulesCore.Tangent{typeof(d)}(; r=r̄)
+        return NoTangent(), d̄, @thunk(project_x(x̄)), @thunk(project_y(ȳ))
+    end
+    return Distances.colwise(d, x, y), colwise_pullback
+end
+
 ## Reverse Rules for matrix wrappers
 
 function ChainRulesCore.rrule(::Type{<:ColVecs}, X::AbstractMatrix)

src/kernels/kernelsum.jl

@@ -45,6 +45,7 @@ Base.length(k::KernelSum) = length(k.kernels)
 
 _sum(f, ks::Tuple, args...) = f(first(ks), args...) + _sum(f, Base.tail(ks), args...)
 _sum(f, ks::Tuple{Tx}, args...) where {Tx} = f(only(ks), args...)
+_sum(f, ks::AbstractVector, args...) = sum(k -> f(k, args...), ks)
 
 (κ::KernelSum)(x, y) = _sum((k, x, y) -> k(x, y), κ.kernels, x, y)
 

test/Project.toml

@@ -1,5 +1,7 @@
 [deps]
 AxisArrays = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"

test/basekernels/periodic.jl

@@ -15,7 +15,6 @@
     TestUtils.test_interface(PeriodicKernel(; r=[0.9, 0.9]), ColVecs{Float64})
     TestUtils.test_interface(PeriodicKernel(; r=[0.8, 0.7]), RowVecs{Float64})
 
-    # test_ADs(r->PeriodicKernel(r =exp.(r)), log.(r), ADs = [:ForwardDiff, :ReverseDiff])
-    @test_broken "Undefined adjoint for Sinus metric, and failing randomly for ForwardDiff and ReverseDiff"
+    test_ADs(r -> PeriodicKernel(; r=exp.(r)), log.(r))
     test_params(k, (r,))
 end

test/basekernels/sm.jl

@@ -57,5 +57,6 @@
     end
 
     # test_ADs(x->spectral_mixture_kernel(exp.(x[1:3]), reshape(x[4:18], 5, 3), reshape(x[19:end], 5, 3)), vcat(log.(αs₁), γs[:], ωs[:]), dims = [5,5])
-    @test_broken "No tests passing (BaseKernel)"
+    # No tests passing (BaseKernel)
+    @test_broken false
 end

test/basekernels/wiener.jl

@@ -44,5 +44,6 @@
     TestUtils.test_interface(k2, x0, x1, x2)
     TestUtils.test_interface(k3, x0, x1, x2)
     # test_ADs(()->WienerKernel(i=1))
-    @test_broken "No tests passing"
+    # No tests passing
+    @test_broken false
 end

test/chainrules.jl

@@ -19,4 +19,23 @@
     compare_gradient(:Zygote, [x, y]) do xy
         KernelFunctions.Sinus(r)(xy[1], xy[2])
     end
+    @testset "rrules for Sinus(r=$r)" for r in (rand(3),)
+        dist = KernelFunctions.Sinus(r)
+        @testset "$type" for type in (Vector, SVector{3})
+            test_rrule(dist, type(rand(3)), type(rand(3)))
+        end
+        @testset "$type1, $type2" for type1 in (Matrix, SMatrix{3,2}),
+            type2 in (Matrix, SMatrix{3,4})
+
+            test_rrule(Distances.pairwise, dist, type1(rand(3, 2)); fkwargs=(dims=2,))
+            test_rrule(
+                Distances.pairwise,
+                dist,
+                type1(rand(3, 2)),
+                type2(rand(3, 4));
+                fkwargs=(dims=2,),
+            )
+            test_rrule(Distances.colwise, dist, type1(rand(3, 2)), type1(rand(3, 2)))
+        end
+    end
 end

test/kernels/kernelsum.jl

@@ -2,21 +2,24 @@
     k1 = LinearKernel()
     k2 = SqExponentialKernel()
     k = KernelSum(k1, k2)
-    @test k == KernelSum([k1, k2]) == KernelSum((k1, k2))
+    kvec = KernelSum([k1, k2])
+    @test k == kvec == KernelSum((k1, k2))
     for (_k1, _k2) in Iterators.product(
         (k1, KernelSum((k1,)), KernelSum([k1])), (k2, KernelSum((k2,)), KernelSum([k2]))
     )
         @test k == _k1 + _k2
+        @test kvec == _k1 + _k2
     end
-    @test length(k) == 2
-    @test repr(k) == (
+    @test length(k) == length(kvec) == 2
+    @test repr(k) ==
+        repr(kvec) ==
         "Sum of 2 kernels:\n" *
         "\tLinear Kernel (c = 0.0)\n" *
         "\tSquared Exponential Kernel (metric = Euclidean(0.0))"
-    )
 
     # Standardised tests.
     test_interface(k, Float64)
+    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

test/runtests.jl

@@ -1,5 +1,7 @@
 using KernelFunctions
 using AxisArrays
+using ChainRulesCore
+using ChainRulesTestUtils
 using Distances
 using Documenter
 using Functors: functor

test/test_utils.jl

@@ -384,7 +384,7 @@ function test_zygote_perf_heuristic(
             @test_broken fwd[1] == fwd[2]
         end
         if passes[3]
-            @test pb[1] == pb[2]
+            @test abs(pb[1] - pb[2]) ≤ 1
         else
             @test_broken pb[1] == pb[2]
         end

test/transform/selecttransform.jl

@@ -104,17 +104,41 @@
     end
 
     @testset "$(AD)" for AD in [:ReverseDiff]
-        @test_broken ga = gradient(AD, A) do a
-            testfunction(ta_row, a, 2)
+        @test_broken let
+            gx = gradient(AD, X) do x
+                testfunction(tx_row, x, 2)
+            end
+            ga = gradient(AD, A) do a
+                testfunction(ta_row, a, 2)
+            end
+            gx ≈ ga
         end
-        @test_broken ga = gradient(AD, A) do a
-            testfunction(ta_col, a, 1)
+        @test_broken let
+            gx = gradient(AD, X) do x
+                testfunction(tx_col, x, 1)
+            end
+            ga = gradient(AD, A) do a
+                testfunction(ta_col, a, 1)
+            end
+            gx ≈ ga
         end
-        @test_broken ga = gradient(AD, A) do a
-            testfunction(ta_row, a, B, 2)
+        @test_broken let
+            gx = gradient(AD, X) do x
+                testfunction(tx_row, x, Y, 2)
+            end
+            ga = gradient(AD, A) do a
+                testfunction(ta_row, a, B, 2)
+            end
+            gx ≈ ga
         end
-        @test_broken ga = gradient(AD, A) do a
-            testfunction(ta_col, a, C, 1)
+        @test_broken let
+            gx = gradient(AD, X) do x
+                testfunction(tx_col, x, Z, 1)
+            end
+            ga = gradient(AD, A) do a
+                testfunction(ta_col, a, C, 1)
+            end
+            gx ≈ ga
         end
     end