Implement Zygote hack (#376)

* Implement Zygote hack

* Update src/KernelFunctions.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update src/utils.jl

Co-authored-by: David Widmann <[email protected]>

* Further tests

* Update test/test_utils.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update test/test_utils.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update test/test_utils.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update test/test_utils.jl

* Update test/utils.jl

* Bump patch version

* Restrict testing to 1.6

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: David Widmann <[email protected]>

Author: 3 people

Project.toml

@@ -1,6 +1,6 @@
 name = "KernelFunctions"
 uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
-version = "0.10.21"
+version = "0.10.22"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/KernelFunctions.jl

@@ -58,7 +58,10 @@ using IrrationalConstants: logtwo, twoπ, invsqrt2
 using LogExpFunctions: softplus
 using StatsBase
 using TensorCore
-using ZygoteRules: ZygoteRules
+using ZygoteRules: ZygoteRules, AContext, literal_getproperty, literal_getfield
+
+# Hack to work around Zygote type inference problems.
+const Distances_pairwise = Distances.pairwise
 
 abstract type Kernel end
 abstract type SimpleKernel <: Kernel end

src/distances/pairwise.jl

@@ -13,11 +13,11 @@ end
 pairwise!(out::AbstractMatrix, d::PreMetric, X::AbstractVector) = pairwise!(out, d, X, X)
 
 function pairwise(d::PreMetric, x::AbstractVector{<:Real})
-    return Distances.pairwise(d, reshape(x, :, 1); dims=1)
+    return Distances_pairwise(d, reshape(x, :, 1); dims=1)
 end
 
 function pairwise(d::PreMetric, x::AbstractVector{<:Real}, y::AbstractVector{<:Real})
-    return Distances.pairwise(d, reshape(x, :, 1), reshape(y, :, 1); dims=1)
+    return Distances_pairwise(d, reshape(x, :, 1), reshape(y, :, 1); dims=1)
 end
 
 function pairwise!(out::AbstractMatrix, d::PreMetric, x::AbstractVector{<:Real})

src/utils.jl

@@ -80,13 +80,13 @@ Base.vcat(a::ColVecs, b::ColVecs) = ColVecs(hcat(a.X, b.X))
 
 dim(x::ColVecs) = size(x.X, 1)
 
-pairwise(d::PreMetric, x::ColVecs) = Distances.pairwise(d, x.X; dims=2)
-pairwise(d::PreMetric, x::ColVecs, y::ColVecs) = Distances.pairwise(d, x.X, y.X; dims=2)
+pairwise(d::PreMetric, x::ColVecs) = Distances_pairwise(d, x.X; dims=2)
+pairwise(d::PreMetric, x::ColVecs, y::ColVecs) = Distances_pairwise(d, x.X, y.X; dims=2)
 function pairwise(d::PreMetric, x::AbstractVector, y::ColVecs)
-    return Distances.pairwise(d, reduce(hcat, x), y.X; dims=2)
+    return Distances_pairwise(d, reduce(hcat, x), y.X; dims=2)
 end
 function pairwise(d::PreMetric, x::ColVecs, y::AbstractVector)
-    return Distances.pairwise(d, x.X, reduce(hcat, y); dims=2)
+    return Distances_pairwise(d, x.X, reduce(hcat, y); dims=2)
 end
 function pairwise!(out::AbstractMatrix, d::PreMetric, x::ColVecs)
     return Distances.pairwise!(out, d, x.X; dims=2)
@@ -150,13 +150,13 @@ Base.vcat(a::RowVecs, b::RowVecs) = RowVecs(vcat(a.X, b.X))
 
 dim(x::RowVecs) = size(x.X, 2)
 
-pairwise(d::PreMetric, x::RowVecs) = Distances.pairwise(d, x.X; dims=1)
-pairwise(d::PreMetric, x::RowVecs, y::RowVecs) = Distances.pairwise(d, x.X, y.X; dims=1)
+pairwise(d::PreMetric, x::RowVecs) = Distances_pairwise(d, x.X; dims=1)
+pairwise(d::PreMetric, x::RowVecs, y::RowVecs) = Distances_pairwise(d, x.X, y.X; dims=1)
 function pairwise(d::PreMetric, x::AbstractVector, y::RowVecs)
-    return Distances.pairwise(d, permutedims(reduce(hcat, x)), y.X; dims=1)
+    return Distances_pairwise(d, permutedims(reduce(hcat, x)), y.X; dims=1)
 end
 function pairwise(d::PreMetric, x::RowVecs, y::AbstractVector)
-    return Distances.pairwise(d, x.X, permutedims(reduce(hcat, y)); dims=1)
+    return Distances_pairwise(d, x.X, permutedims(reduce(hcat, y)); dims=1)
 end
 function pairwise!(out::AbstractMatrix, d::PreMetric, x::RowVecs)
     return Distances.pairwise!(out, d, x.X; dims=1)

src/zygoterules.jl

@@ -5,3 +5,9 @@ end
 ZygoteRules.@adjoint function Base.map(t::Transform, X::RowVecs)
     return ZygoteRules.pullback(_map, t, X)
 end
+
+function ZygoteRules._pullback(
+    cx::AContext, ::typeof(literal_getproperty), x::ColVecs, ::Val{f}
+) where {f}
+    return ZygoteRules._pullback(cx, literal_getfield, x, Val{f}())
+end

test/test_utils.jl

@@ -31,6 +31,15 @@ end
 
 # AD utilities
 
+# Type to work around some performance issues that can happen on the reverse-pass of Zygote.
+# This context doesn't allow any globals. Don't use this if you use globals in your
+# programme.
+struct NoContext <: Zygote.AContext end
+
+Zygote.cache(cx::NoContext) = (cache_fields = nothing)
+Base.haskey(cx::NoContext, x) = false
+Zygote.accum_param(::NoContext, x, Δ) = Δ
+
 const FDM = FiniteDifferences.central_fdm(5, 1)
 
 gradient(f, s::Symbol, args) = gradient(f, Val(s), args)
@@ -87,6 +96,13 @@ function test_ADs(
     end
 end
 
+function check_zygote_type_stability(f, args...; ctx=Zygote.Context())
+    @inferred f(args...)
+    @inferred Zygote._pullback(ctx, f, args...)
+    out, pb = Zygote._pullback(ctx, f, args...)
+    @inferred pb(out)
+end
+
 function test_ADs(
     k::MOKernel; ADs=[:Zygote, :ForwardDiff, :ReverseDiff], dims=(in=3, out=2, obs=3)
 )

test/utils.jl

@@ -51,6 +51,22 @@
             X_, back = Zygote.pullback(DX -> DX.X, DX)
             @test back(ones(size(X)))[1].X == ones(size(X))
         end
+
+        if VERSION >= v"1.6"
+            @testset "Zygote type-inference" begin
+                ctx = NoContext()
+                x = ColVecs(randn(2, 4))
+                y = ColVecs(randn(2, 3))
+
+                # Ensure KernelFunctions.pairwise rather than Distances.pairwise is used.
+                check_zygote_type_stability(
+                    x -> KernelFunctions.pairwise(SqEuclidean(), x), x; ctx=ctx
+                )
+                check_zygote_type_stability(
+                    (x, y) -> KernelFunctions.pairwise(SqEuclidean(), x, y), x, y; ctx=ctx
+                )
+            end
+        end
     end
     @testset "RowVecs" begin
         DX = RowVecs(X)