Fix gradient issues with kernelmatrix_diag and use ChainRulesCore

Project.toml

@@ -3,6 +3,7 @@ uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
 version = "0.8.24"
 
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
@@ -17,6 +18,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
+ChainRulesCore = "0.9"
 Compat = "3.7"
 Distances = "0.9.1, 0.10"
 Functors = "0.1"

src/KernelFunctions.jl

@@ -55,11 +55,12 @@ export IndependentMOKernel, LatentFactorMOKernel
 export tensor, ⊗
 
 using Compat
+using ChainRulesCore
 using Requires
 using Distances, LinearAlgebra
 using Functors
 using SpecialFunctions: loggamma, besselk, polygamma
-using ZygoteRules: @adjoint, pullback
+# using ZygoteRules: @adjoint, pullback, ZygoteRules
 using StatsFuns: logtwo
 using StatsBase
 using TensorCore
@@ -112,7 +113,7 @@ include(joinpath("mokernels", "moinput.jl"))
 include(joinpath("mokernels", "independent.jl"))
 include(joinpath("mokernels", "slfm.jl"))
 
-include("zygote_adjoints.jl")
+include("chainrules.jl")
 
 include("test_utils.jl")
 

src/chainrules.jl

@@ -0,0 +1,142 @@
+## Reverse Rules Delta
+
+function rrule(::typeof(Distances.evaluate), s::Delta, x::AbstractVector, y::AbstractVector)
+    d = evaluate(s, x, y)
+    function evaluate_pullback(::Any)
+        return NO_FIELDS, Zero(), Zero()
+    end
+    return d, evaluate_pullback
+end
+
+function rrule(
+    ::typeof(Distances.pairwise), d::Delta, X::AbstractMatrix, Y::AbstractMatrix; dims=2
+)
+    P = Distances.pairwise(d, X, Y; dims=dims)
+    function pairwise_pullback(::Any)
+        return NO_FIELDS, Zero(), Zero()
+    end
+    return P, pairwise_pullback
+end
+
+function rrule(::typeof(Distances.pairwise), d::Delta, X::AbstractMatrix; dims=2)
+    P = Distances.pairwise(d, X; dims=dims)
+    function pairwise_pullback(::Any)
+        return NO_FIELDS, Zero()
+    end
+    return P, pairwise_pullback
+end
+
+function rrule(::typeof(Distances.colwise), d::Delta, X::AbstractMatrix, Y::AbstractMatrix)
+    C = Distances.colwise(d, X, Y)
+    function colwise_pullback(::AbstractVector)
+        return NO_FIELDS, Zero(), Zero()
+    end
+    return C, colwise_pullback
+end
+
+## Reverse Rules DotProduct
+function rrule(
+    ::typeof(Distances.evaluate), s::DotProduct, x::AbstractVector, y::AbstractVector
+)
+    d = dot(x, y)
+    function evaluate_pullback(Δ)
+        return NO_FIELDS, Δ .* y, Δ .* x
+    end
+    return d, evaluate_pullback
+end
+
+function rrule(
+    ::typeof(Distances.pairwise),
+    d::DotProduct,
+    X::AbstractMatrix,
+    Y::AbstractMatrix;
+    dims=2,
+)
+    P = Distances.pairwise(d, X, Y; dims=dims)
+    if dims == 1
+        function pairwise_pullback_cols(Δ)
+            return NO_FIELDS, Δ * Y, Δ' * X
+        end
+        return P, pairwise_pullback_cols
+    else
+        function pairwise_pullback_rows(Δ)
+            return NO_FIELDS, Y * Δ', X * Δ
+        end
+        return P, pairwise_pullback_rows
+    end
+end
+
+function rrule(::typeof(Distances.pairwise), d::DotProduct, X::AbstractMatrix; dims=2)
+    P = Distances.pairwise(d, X; dims=dims)
+    if dims == 1
+        function pairwise_pullback_cols(Δ)
+            return NO_FIELDS, 2 * Δ * X
+        end
+        return P, pairwise_pullback_cols
+    else
+        function pairwise_pullback_rows(Δ)
+            return NO_FIELDS, 2 * X * Δ
+        end
+        return P, pairwise_pullback_rows
+    end
+end
+
+function rrule(
+    ::typeof(Distances.colwise), d::DotProduct, X::AbstractMatrix, Y::AbstractMatrix
+)
+    C = Distances.colwise(d, X, Y)
+    function colwise_pullback(Δ::AbstractVector)
+        return (nothing, Δ' .* Y, Δ' .* X)
+    end
+    return C, colwise_pullback
+end
+
+## Reverse Rules Sinus
+function rrule(::typeof(Distances.evaluate), s::Sinus, x::AbstractVector, y::AbstractVector)
+    d = (x - y)
+    sind = sinpi.(d)
+    val = sum(abs2, sind ./ s.r)
+    gradx = 2π .* cospi.(d) .* sind ./ (s.r .^ 2)
+    function evaluate_pullback(Δ)
+        return (r=-2Δ .* abs2.(sind) ./ s.r,), Δ * gradx, -Δ * gradx
+    end
+    return val, evaluate_pullback
+end
+
+## Reverse Rules for matrix wrappers
+
+function rrule(::ColVecs, X::AbstractMatrix)
+    ColVecs_pullback(Δ::NamedTuple) = (Δ.X,)
+    ColVecs_pullback(Δ::AbstractMatrix) = (Δ,)
+    function ColVecs_pullback(Δ::AbstractVector{<:AbstractVector{<:Real}})
+        return throw(error("In slow method"))
+    end
+    return ColVecs(X), ColVecs_pullback
+end
+
+function rrule(::RowVecs, X::AbstractMatrix)
+    RowVecs_pullback(Δ::NamedTuple) = (Δ.X,)
+    RowVecs_pullback(Δ::AbstractMatrix) = (Δ,)
+    function RowVecs_pullback(Δ::AbstractVector{<:AbstractVector{<:Real}})
+        return throw(error("In slow method"))
+    end
+    return RowVecs(X), RowVecs_pullback
+end
+
+# function rrule(::typeof(Base.map), t::Transform, X::ColVecs)
+#     return pullback(_map, t, X)
+# end
+
+# function rrule(::typeof(Base.map), t::Transform, X::RowVecs)
+#     return pullback(_map, t, X)
+# end
+
+# @adjoint function (dist::Distances.SqMahalanobis)(a, b)
+#     function SqMahalanobis_pullback(Δ::Real)
+#         B_Bᵀ = dist.qmat + transpose(dist.qmat)
+#         a_b = a - b
+#         δa = (B_Bᵀ * a_b) * Δ
+#         return (qmat=(a_b * a_b') * Δ,), δa, -δa
+#     end
+#     return evaluate(dist, a, b), SqMahalanobis_pullback
+# end

src/distances/pairwise.jl

@@ -29,3 +29,29 @@ function pairwise!(
 )
     return Distances.pairwise!(out, d, reshape(x, :, 1), reshape(y, :, 1); dims=1)
 end
+
+# Also defines the colwise method for abstractvectors
+
+function colwise(d::PreMetric, x::ColVecs)
+    return Distances.colwise(d, x.X, x.X)
+end
+
+function colwise(d::PreMetric, x::RowVecs)
+    return Distances.colwise(d, x.X', x.X')
+end
+
+function colwise(d::PreMetric, x::AbstractVector)
+    return map(d, x, x)
+end
+
+function colwise(d::PreMetric, x::ColVecs, y::ColVecs)
+    return Distances.colwise(d, x.X, y.X)
+end
+
+function colwise(d::PreMetric, x::RowVecs, y::RowVecs)
+    return Distances.colwise(d, x.X', y.X')
+end
+
+function colwise(d::PreMetric, x::AbstractVector, y::AbstractVector)
+    return map(d, x, y)
+end

src/kernels/transformedkernel.jl

@@ -80,6 +80,12 @@ function kernelmatrix_diag!(K::AbstractVector, κ::TransformedKernel, x::Abstrac
     return kernelmatrix_diag!(K, κ.kernel, _map(κ.transform, x))
 end
 
+function kernelmatrix_diag!(
+    K::AbstractVector, κ::TransformedKernel, x::AbstractVector, y::AbstractVector
+)
+    return kernelmatrix_diag!(K, κ.kernel, _map(κ.transform, x), _map(κ.transform, y))
+end
+
 function kernelmatrix!(K::AbstractMatrix, κ::TransformedKernel, x::AbstractVector)
     return kernelmatrix!(K, kernel(κ), _map(κ.transform, x))
 end
@@ -94,6 +100,10 @@ function kernelmatrix_diag(κ::TransformedKernel, x::AbstractVector)
     return kernelmatrix_diag(κ.kernel, _map(κ.transform, x))
 end
 
+function kernelmatrix_diag(κ::TransformedKernel, x::AbstractVector, y::AbstractVector)
+    return kernelmatrix_diag(κ.kernel, _map(κ.transform, x), _map(κ.transform, y))
+end
+
 function kernelmatrix(κ::TransformedKernel, x::AbstractVector)
     return kernelmatrix(kernel(κ), _map(κ.transform, x))
 end

src/matrix/kernelmatrix.jl

@@ -80,24 +80,32 @@ kernelmatrix_diag(κ::Kernel, x::AbstractVector, y::AbstractVector) = map(κ, x,
 function kernelmatrix!(K::AbstractMatrix, κ::SimpleKernel, x::AbstractVector)
     validate_inplace_dims(K, x)
     pairwise!(K, metric(κ), x)
-    return map!(d -> kappa(κ, d), K, K)
+    return map!(Base.Fix1(kappa, κ), K, K)
 end
 
 function kernelmatrix!(
     K::AbstractMatrix, κ::SimpleKernel, x::AbstractVector, y::AbstractVector
 )
     validate_inplace_dims(K, x, y)
     pairwise!(K, metric(κ), x, y)
-    return map!(d -> kappa(κ, d), K, K)
+    return map!(Base.Fix1(kappa, κ), K, K)
 end
 
 function kernelmatrix(κ::SimpleKernel, x::AbstractVector)
-    return map(d -> kappa(κ, d), pairwise(metric(κ), x))
+    return map(Base.Fix1(kappa, κ), pairwise(metric(κ), x))
 end
 
 function kernelmatrix(κ::SimpleKernel, x::AbstractVector, y::AbstractVector)
     validate_inputs(x, y)
-    return map(d -> kappa(κ, d), pairwise(metric(κ), x, y))
+    return map(Base.Fix1(kappa, κ), pairwise(metric(κ), x, y))
+end
+
+function kernelmatrix_diag(κ::SimpleKernel, x::AbstractVector)
+    return map(Base.Fix1(kappa, κ), colwise(metric(κ), x))
+end
+
+function kernelmatrix_diag(κ::SimpleKernel, x::AbstractVector, y::AbstractVector)
+    return map(Base.Fix1(kappa, κ), colwise(metric(κ), x, y))
 end
 
 #