Fix AD issue with mixed input for TransformedKernel (#160)

* Fix mixed input bug for TransformedKernel

* Fix tests

* Support mized input in pairwise

* avoid splatting

* style fixes

* use permute dims

Author: sharanry

src/utils.jl

@@ -43,6 +43,12 @@ 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)
+function pairwise(d::PreMetric, x::AbstractVector, y::ColVecs)
+    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)
+end
 function pairwise!(out::AbstractMatrix, d::PreMetric, x::ColVecs)
     return Distances.pairwise!(out, d, x.X; dims=2)
 end
@@ -73,6 +79,12 @@ 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)
+function pairwise(d::PreMetric, x::AbstractVector, y::RowVecs)
+    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)
+end
 function pairwise!(out::AbstractMatrix, d::PreMetric, x::RowVecs)
     return Distances.pairwise!(out, d, x.X; dims=1)
 end

test/kernels/transformedkernel.jl

@@ -50,6 +50,22 @@
             tmp_diag = Vector{Float64}(undef, length(x))
             @test kerneldiagmatrix!(tmp_diag, kt, x) ≈ kerneldiagmatrix(kt, x)
         end
+
+        @testset "mixed inputs" begin
+            k = transform(SqExponentialKernel(), 10.0)
+            x = rand(rng, 5, 3)
+            X1 = collect(eachcol(x))
+            Y1 = KernelFunctions.ColVecs(x)
+            @test_nowarn Zygote.gradient(k-> sum(kernelmatrix(k, X1, Y1)), k)
+            @test_nowarn Zygote.gradient(k-> sum(kernelmatrix(k, Y1, X1)), k)
+            @test kernelmatrix(k, X1, Y1) ≈ kernelmatrix(k, X1, X1) ≈ kernelmatrix(k, Y1, Y1)
+
+            X2 = collect(eachrow(x))
+            Y2 = KernelFunctions.RowVecs(x)
+            @test_nowarn Zygote.gradient(k-> sum(kernelmatrix(k, X2, Y2)), k)
+            @test_nowarn Zygote.gradient(k-> sum(kernelmatrix(k, Y2, X2)), k)
+            @test kernelmatrix(k, X2, Y2) ≈ kernelmatrix(k, X2, X2) ≈ kernelmatrix(k, Y2, Y2)
+        end
     end
     test_ADs(x->transform(SqExponentialKernel(), x[1]), rand(1))# ADs = [:ForwardDiff, :ReverseDiff])
     # Test implicit gradients