Merge pull request #15 from theogf/fixdiagmatrix

Added adjoints for evaluate

Author: theogf

.travis.yml

@@ -5,7 +5,13 @@ os:
   - osx
 julia:
   - 1.0
+  - 1.2
+  - 1.3
   - nightly
+# because of Zygote needs to allow failing on nightly
+matrix:
+  allow_failures:
+    - julia: nightly
 notifications:
   email: false
 after_success:

Project.toml

@@ -3,22 +3,27 @@ uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
 version = "0.2.2"
 
 [deps]
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
+Compat = "2.2"
 Distances = "0.8"
 PDMats = "0.9"
 SpecialFunctions = "0"
 StatsFuns = "0.8"
+Zygote = "0.4"
 julia = "1.0"
 
 [extras]
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 
 [targets]
-test = ["Random", "Test"]
+test = ["Random", "Test", "FiniteDifferences"]

src/KernelFunctions.jl

@@ -14,8 +14,10 @@ export KernelSum, KernelProduct
 
 export Transform, SelectTransform, ChainTransform, ScaleTransform, LowRankTransform, IdentityTransform, FunctionTransform
 
+using Compat
 using Distances, LinearAlgebra
 using SpecialFunctions: lgamma, besselk
+using Zygote: @adjoint
 using StatsFuns: logtwo
 using PDMats: PDMat
 
@@ -42,4 +44,6 @@ include("kernels/kernelproduct.jl")
 
 include("generic.jl")
 
+include("zygote_adjoints.jl")
+
 end

src/matrix/kernelmatrix.jl

@@ -108,9 +108,9 @@ function kerneldiagmatrix(
         )
         @assert obsdim ∈ [1,2] "obsdim should be 1 or 2 (see docs of kernelmatrix))"
         if obsdim == 1
-            [@views _kernel(κ,X[i,:],X[i,:]) for i in 1:size(X,obsdim)]
+            @compat eachrow(X) .|> x->_kernel(κ,x,x) #[@views _kernel(κ,X[i,:],X[i,:]) for i in 1:size(X,obsdim)]
         elseif obsdim == 2
-            [@views _kernel(κ,X[:,i],X[:,i]) for i in 1:size(X,obsdim)]
+            @compat eachcol(X) .|> x->_kernel(κ,x,x) #[@views _kernel(κ,X[:,i],X[:,i]) for i in 1:size(X,obsdim)]
         end
 end
 

src/zygote_adjoints.jl

@@ -0,0 +1,22 @@
+@adjoint function evaluate(s::SqEuclidean, x::AbstractVector, y::AbstractVector)
+  δ = x .- y
+  sum(abs2, δ), Δ -> begin
+    x̄ = (2 * Δ) .* δ
+    (nothing, x̄, -x̄)
+  end
+end
+
+@adjoint function evaluate(s::Euclidean, x::AbstractVector, y::AbstractVector)
+  D = x.-y
+  δ = sqrt(sum(abs2,D))
+  δ, Δ -> begin
+    x̄ = Δ .* D / (δ + eps(δ))
+    (nothing, x̄, -x̄)
+  end
+end
+
+@adjoint function evaluate(s::DotProduct, x::AbstractVector, y::AbstractVector)
+  dot(x,y), Δ -> begin
+    (nothing, Δ.*y, Δ.*x)
+  end
+end

test/runtests.jl

@@ -11,5 +11,6 @@ include("test_transform.jl")
 include("test_distances.jl")
 include("test_kernels.jl")
 include("test_generic.jl")
+include("test_adjoints.jl")
     #include("types.jl")
 end

test/test_adjoints.jl

@@ -0,0 +1,25 @@
+using Zygote
+using FiniteDifferences
+using KernelFunctions
+using Distances
+using Test
+
+x = rand(5)
+y = rand(5)
+
+@testset "Testing Zygote adjoints" begin
+gzeucl = first(Zygote.gradient(xy->evaluate(Euclidean(),xy[1],xy[2]),[x,y]))
+gzsqeucl =  first(Zygote.gradient(xy->evaluate(SqEuclidean(),xy[1],xy[2]),[x,y]))
+gzdotprod = first(Zygote.gradient(xy->evaluate(KernelFunctions.DotProduct(),xy[1],xy[2]),[x,y]))
+
+FDM = central_fdm(5,1)
+
+gfeucl = collect(first(FiniteDifferences.grad(FDM,xy->evaluate(Euclidean(),xy[1],xy[2]),(x,y))))
+gfsqeucl = collect(first(FiniteDifferences.grad(FDM,xy->evaluate(SqEuclidean(),xy[1],xy[2]),(x,y))))
+gfdotprod =collect(first(FiniteDifferences.grad(FDM,xy->evaluate(KernelFunctions.DotProduct(),xy[1],xy[2]),(x,y))))
+
+@test all(gzeucl .≈ gfeucl)
+@test all(gzsqeucl .≈ gfsqeucl)
+@test all(gzdotprod .≈ gfdotprod)
+
+end