Merge pull request #5 from theogf/wct/zygote-euclidean

Add Euclidean rules + refactor Project.toml

Author: theogf

Project.toml

@@ -5,4 +5,15 @@ version = "0.1.0"
 [deps]
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[compat]
+FiniteDifferences = ">= 0.7.2"
+
+[extras]
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["FiniteDifferences", "Random", "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 Zygote: @adjoint, forward
+
+@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 = 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 = forward(X->pairwise(SqEuclidean(), X; dims=2), X)
+    D .= sqrt.(D)
+    return D, function(Δ)
+        Δ = Δ ./ (2 .* D)
+        Δ[diagind(Δ)] .= 0
+        return (nothing, first(back(Δ)))
+    end
+end

test/runtests.jl

@@ -1,6 +1,15 @@
 using Test
 using KernelFunctions
 using Distances
+using FiniteDifferences
+using Random
+using Zygote
 
-include("kernelmatrix.jl")
-include("constructors.jl")
+# Helpful functionality for writing tests.
+include("test_util.jl")
+
+@testset "KernelFunctions" begin
+    include("zygote_rules.jl")
+    include("kernelmatrix.jl")
+    include("constructors.jl")    
+end

test/test_util.jl

@@ -0,0 +1,79 @@
+using FiniteDifferences: j′vp
+
+# Default tolerances for testing Zygote against FiniteDifferences.
+const _rtol = 1e-10
+const _atol = 1e-10
+
+
+
+#
+# Print stuff for debugging
+#
+
+function print_adjoints(adjoint_ad, adjoint_fd, rtol, atol)
+    @show typeof(adjoint_ad), typeof(adjoint_fd)
+    adjoint_ad, adjoint_fd = to_vec(adjoint_ad)[1], to_vec(adjoint_fd)[1]
+    println("atol is $atol, rtol is $rtol")
+    println("ad, fd, abs, rel")
+    abs_err = abs.(adjoint_ad .- adjoint_fd)
+    rel_err = abs.((adjoint_ad .- adjoint_fd) ./ adjoint_ad)
+    display([adjoint_ad adjoint_fd abs_err rel_err])
+    println()
+end
+
+
+
+#
+# Version of isapprox that works for lots of types.
+#
+
+function fd_isapprox(x_ad::Nothing, x_fd, rtol, atol)
+    return fd_isapprox(x_fd, zero(x_fd), rtol, atol)
+end
+function fd_isapprox(x_ad::AbstractArray, x_fd::AbstractArray, rtol, atol)
+    return all(fd_isapprox.(x_ad, x_fd, rtol, atol))
+end
+function fd_isapprox(x_ad::Real, x_fd::Real, rtol, atol)
+    return isapprox(x_ad, x_fd; rtol=rtol, atol=atol)
+end
+function fd_isapprox(x_ad::NamedTuple, x_fd, rtol, atol)
+    f = (x_ad, x_fd)->fd_isapprox(x_ad, x_fd, rtol, atol)
+    return all([f(getfield(x_ad, key), getfield(x_fd, key)) for key in keys(x_ad)])
+end
+function fd_isapprox(x_ad::Tuple, x_fd::Tuple, rtol, atol)
+    return all(map((x, x′)->fd_isapprox(x, x′, rtol, atol), x_ad, x_fd))
+end
+function fd_isapprox(x_ad::Dict, x_fd::Dict, rtol, atol)
+    return all([fd_isapprox(get(()->nothing, x_ad, key), x_fd[key], rtol, atol) for
+        key in keys(x_fd)])
+end
+
+
+
+#
+# Check Zygote against FiniteDifferences.
+#
+
+function adjoint_test(
+    f, ȳ, x...;
+    rtol=_rtol,
+    atol=_atol,
+    fdm=FiniteDifferences.Central(5, 1),
+    print_results=false,
+)
+
+    # Compute forwards-pass and j′vp.
+    y, back = Zygote.forward(f, x...)
+    adj_ad = back(ȳ)
+    adj_fd = j′vp(fdm, f, ȳ, x...)
+
+    # If unary, pull out first thing from ad.
+    adj_ad = length(x) == 1 ? first(adj_ad) : adj_ad
+
+    # Check that forwards-pass agrees with plain forwards-pass.
+    @test y ≈ f(x...)
+
+    # Check that ad and fd adjoints (approximately) agree.
+    print_results && print_adjoints(adj_ad, adj_fd, rtol, atol)
+    @test fd_isapprox(adj_ad, adj_fd, rtol, atol)
+end

test/zygote_rules.jl

@@ -0,0 +1,17 @@
+@testset "zygote_rules" begin
+    @testset "colwise(::Euclidean, X, Y; dims=2)" begin
+        rng, D, P = MersenneTwister(123456), 2, 3
+        X, Y, D̄ = randn(rng, D, P), randn(rng, D, P), randn(rng, P)
+        adjoint_test((X, Y)->colwise(Euclidean(), X, Y), D̄, X, Y)
+    end
+    @testset "pairwise(::Euclidean, X, Y; dims=2)" begin
+        rng, D, P, Q = MersenneTwister(123456), 2, 3, 5
+        X, Y, D̄ = randn(rng, D, P), randn(rng, D, Q), randn(rng, P, Q)
+        adjoint_test((X, Y)->pairwise(Euclidean(), X, Y; dims=2), D̄, X, Y)
+    end
+    @testset "pairwise(::Euclidean, X; dims=2)" begin
+        rng, D, P = MersenneTwister(123456), 2, 3
+        X, D̄ = randn(rng, D, P), randn(rng, P, P)
+        adjoint_test(X->pairwise(Euclidean(), X; dims=2), D̄, X)
+    end
+end