Fix periodic kernel AD

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

@@ -112,15 +112,118 @@ end
 function ChainRulesCore.rrule(s::Sinus, x::AbstractVector, y::AbstractVector)
     d = x - y
     sind = sinpi.(d)
-    abs2_sind_r = abs2.(sind) ./ s.r
-    val = sum(abs2_sind_r)
-    gradx = twoπ .* cospi.(d) .* sind ./ (s.r .^ 2)
+    abs2_sind_r = abs2.(sind) ./ s.r .^ 2
+    val = sum(abs2_sind_r)  
+    gradx = twoπ .* cospi.(d) .* sind ./ s.r .^ 2
     function evaluate_pullback(Δ::Any)
-        return (r=-2Δ .* abs2_sind_r,), Δ * gradx, -Δ * gradx
+        return (r=-2Δ .* abs2_sind_r ./ s.r,), Δ * 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̄ = 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 = 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
+        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
+        NoTangent(), (r=r̄,), @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̄ = zero(x)
+        ȳ = 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 = twoπ .* Δ[i, j] .* sinpi.(xi .- yj) .* cospi.(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 = twoπ .* Δ[i, j] .* sinpi.(xi .- yj) .* cospi.(xi .- yj) ./ d.r .^ 2
+                r̄ .-= 2 .* Δ[i, j] .* sinpi.(xi .- yj) .^ 2 ./ d.r .^ 3
+                x̄[:, i] += ds
+                ȳ[:, j] -= ds
+            end
+        end
+        NoTangent(), (r=r̄,), @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̄ = zero(x)
+        ȳ = zero(y)
+        r̄ = zero(d.r)
+        for i in 1:n
+            xi = view(x, :, i)
+            yi = view(y, :, i)
+            ds = twoπ .* Δ[i] .* sinpi.(xi .- yi) .* cospi.(xi .- yi) ./ d.r .^ 2
+            r̄ .-= 2 .* Δ[i] .* sinpi.(xi .- yi) .^ 2 ./ d.r .^ 3
+            x̄[:, i] += ds
+            ȳ[:, i] -= ds
+        end
+        NoTangent(), (r=r̄,), @thunk(project_x(x̄)), @thunk(project_y(ȳ))
+    end
+    return Distances.colwise(d, x, y), colwise_pullback
+end
+
 ## Reverse Rules SqMahalanobis
 
 function ChainRulesCore.rrule(

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