Add basic tests for rrules

Author: simsurace

src/chainrules.jl

@@ -116,7 +116,9 @@ function ChainRulesCore.rrule(s::Sinus, x::AbstractVector, y::AbstractVector)
     val = sum(abs2_sind_r)
     gradx = twoπ .* cospi.(d) .* sind ./ s.r .^ 2
     function evaluate_pullback(Δ::Any)
-        return (r=-2Δ .* abs2_sind_r ./ s.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
@@ -149,7 +151,8 @@ function ChainRulesCore.rrule(
                 x̄[:, j] -= ds
             end
         end
-        return NoTangent(), (r=r̄,), @thunk(project_x(x̄))
+        d̄ = ChainRulesCore.Tangent{typeof(d)}(; r=r̄)
+        return NoTangent(), d̄, @thunk(project_x(x̄))
     end
     return Distances.pairwise(d, x; dims), pairwise_pullback
 end
@@ -185,7 +188,8 @@ function ChainRulesCore.rrule(
                 ȳ[:, j] -= ds
             end
         end
-        return NoTangent(), (r=r̄,), @thunk(project_x(x̄)), @thunk(project_y(ȳ))
+        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
@@ -209,7 +213,8 @@ function ChainRulesCore.rrule(
             x̄[:, i] += ds
             ȳ[:, i] -= ds
         end
-        return NoTangent(), (r=r̄,), @thunk(project_x(x̄)), @thunk(project_y(ȳ))
+        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

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/chainrules.jl

@@ -28,4 +28,12 @@
             SqMahalanobis(Qxy[1])(Qxy[2], Qxy[3])
         end
     end
+
+    @testset "rrules for Sinus(r=$r)" for r in (rand(3),)
+        dist = KernelFunctions.Sinus(r)
+        ddist = (r = ones(length(r)),)
+        test_rrule(dist, rand(3), rand(3))
+        test_rrule(Distances.pairwise, dist, rand(3, 2); fkwargs=(dims=2,))
+        test_rrule(Distances.pairwise, dist, rand(3, 2), rand(3, 3); fkwargs=(dims=2,))
+    end
 end

test/runtests.jl

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