Add Eucldiean rules + refactor Project.toml

Author: willtebbutt

Project.toml

@@ -5,4 +5,10 @@ version = "0.1.0"
 [deps]
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
+
+[extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["Test"]

src/KernelFunctions.jl

@@ -8,6 +8,7 @@ using Distances, LinearAlgebra
 const defaultobs = 2
 abstract type Kernel{T<:Real} end
 
+include("zygote_rules.jl")
 include("utils.jl")
 include("common.jl")
 include("kernelmatrix.jl")

src/zygote_rules.jl

@@ -0,0 +1,27 @@
+using ZygoteRules
+
+@adjoint function colwise(s::Euclidean, x::AbstractMatrix, y::AbstractMatrix)
+    d = colwise(s, x, y)
+    return d, function (Δ::AbstractVector)
+        x̄ = (Δ ./ d)' .* (x .- y)
+        return nothing, x̄, -x̄
+    end
+end
+
+@adjoint function pairwise(::Euclidean, X::AbstractMatrix, Y::AbstractMatrix; dims=2)
+    @assert dims == 2
+    D, back = Zygote.forward((X, Y)->pairwise(SqEuclidean(), X, Y; dims=2), X, Y)
+    D .= sqrt.(D)
+    return D, Δ -> (nothing, back(Δ ./ (2 .* D))...)
+end
+
+@adjoint function pairwise(::Euclidean, X::AbstractMatrix; dims=2)
+    @assert dims == 2
+    D, back = Zygote.forward(X->pairwise(SqEuclidean(), X; dims=2), X)
+    D .= sqrt.(D)
+    return D, function(Δ)
+        Δ = Δ ./ (2 .* D)
+        Δ[diagind(Δ)] .= 0
+        return (nothing, first(back(Δ)))
+    end
+end