Add lazy kronecker product for matrix kernels, if Kronecker.jl is loaded (#364)

* Restore additions

* Improvements for lazy kron

* Remove unneeded lines

* Small experiment with overwriting

* Reorder and overwrite

* Format and kernelmatrix!

* Reinstate separate method

* Adding tests

* Duplicate code for readability

* Format

* Remove comment and patch bump

* Change kernelmatrix!

* Change to output covariance type

* Change to output covariance type - revert

* Revert "Change to output covariance type - revert"

This reverts commit 09cd20e.

* Revert "Change kernelmatrix!"

This reverts commit f46bd61.

* Add kernelmatrix! changes again

* Change input types for pairwise pullback

* Missing changes to Any

Author: Crown421

Project.toml

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

src/KernelFunctions.jl

@@ -105,8 +105,8 @@ include("kernels/neuralkernelnetwork.jl")
 include("approximations/nystrom.jl")
 include("generic.jl")
 
-include("mokernels/mokernel.jl")
 include("mokernels/moinput.jl")
+include("mokernels/mokernel.jl")
 include("mokernels/independent.jl")
 include("mokernels/slfm.jl")
 include("mokernels/intrinsiccoregion.jl")

src/chainrules.jl

@@ -26,7 +26,7 @@ function ChainRulesCore.rrule(
     ::typeof(Distances.pairwise), d::Delta, X::AbstractMatrix, Y::AbstractMatrix; dims=2
 )
     P = Distances.pairwise(d, X, Y; dims=dims)
-    function pairwise_pullback(::AbstractMatrix)
+    function pairwise_pullback(::Any)
         return NoTangent(), NoTangent(), ZeroTangent(), ZeroTangent()
     end
     return P, pairwise_pullback
@@ -36,7 +36,7 @@ function ChainRulesCore.rrule(
     ::typeof(Distances.pairwise), d::Delta, X::AbstractMatrix; dims=2
 )
     P = Distances.pairwise(d, X; dims=dims)
-    function pairwise_pullback(::AbstractMatrix)
+    function pairwise_pullback(::Any)
         return NoTangent(), NoTangent(), ZeroTangent()
     end
     return P, pairwise_pullback
@@ -46,7 +46,7 @@ function ChainRulesCore.rrule(
     ::typeof(Distances.colwise), d::Delta, X::AbstractMatrix, Y::AbstractMatrix
 )
     C = Distances.colwise(d, X, Y)
-    function colwise_pullback(::AbstractVector)
+    function colwise_pullback(::Any)
         return NoTangent(), NoTangent(), ZeroTangent(), ZeroTangent()
     end
     return C, colwise_pullback
@@ -70,7 +70,7 @@ function ChainRulesCore.rrule(
     dims=2,
 )
     P = Distances.pairwise(d, X, Y; dims=dims)
-    function pairwise_pullback_cols(Δ::AbstractMatrix)
+    function pairwise_pullback_cols(Δ::Any)
         if dims == 1
             return NoTangent(), NoTangent(), Δ * Y, Δ' * X
         else
@@ -84,7 +84,7 @@ function ChainRulesCore.rrule(
     ::typeof(Distances.pairwise), d::DotProduct, X::AbstractMatrix; dims=2
 )
     P = Distances.pairwise(d, X; dims=dims)
-    function pairwise_pullback_cols(Δ::AbstractMatrix)
+    function pairwise_pullback_cols(Δ::Any)
         if dims == 1
             return NoTangent(), NoTangent(), 2 * Δ * X
         else
@@ -98,7 +98,7 @@ function ChainRulesCore.rrule(
     ::typeof(Distances.colwise), d::DotProduct, X::AbstractMatrix, Y::AbstractMatrix
 )
     C = Distances.colwise(d, X, Y)
-    function colwise_pullback(Δ::AbstractVector)
+    function colwise_pullback(Δ::Any)
         return NoTangent(), NoTangent(), Δ' .* Y, Δ' .* X
     end
     return C, colwise_pullback

src/matrix/kernelkroneckermat.jl

@@ -4,6 +4,7 @@
 using .Kronecker: Kronecker
 
 export kernelkronmat
+export kronecker_kernelmatrix
 
 function kernelkronmat(κ::Kernel, X::AbstractVector, dims::Int)
     @assert iskroncompatible(κ) "The chosen kernel is not compatible for kroenecker matrices (see [`iskroncompatible`](@ref))"
@@ -25,3 +26,42 @@ end
     k(x,x') = ∏ᵢᴰ k(xᵢ,x'ᵢ)
 """
 @inline iskroncompatible(κ::Kernel) = false # Default return for kernels
+
+function _kernelmatrix_kroneckerjl_helper(::MOInputIsotopicByFeatures, Kfeatures, Koutputs)
+    return Kronecker.kronecker(Kfeatures, Koutputs)
+end
+
+function _kernelmatrix_kroneckerjl_helper(::MOInputIsotopicByOutputs, Kfeatures, Koutputs)
+    return Kronecker.kronecker(Koutputs, Kfeatures)
+end
+
+function kronecker_kernelmatrix(
+    k::Union{IndependentMOKernel,IntrinsicCoregionMOKernel},
+    x::IsotopicMOInputsUnion,
+    y::IsotopicMOInputsUnion,
+)
+    @assert x.out_dim == y.out_dim
+    Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
+    Koutputs = _mo_output_covariance(k, x.out_dim)
+    return _kernelmatrix_kroneckerjl_helper(x, Kfeatures, Koutputs)
+end
+
+function kronecker_kernelmatrix(
+    k::Union{IndependentMOKernel,IntrinsicCoregionMOKernel}, x::IsotopicMOInputsUnion
+)
+    Kfeatures = kernelmatrix(k.kernel, x.x)
+    Koutputs = _mo_output_covariance(k, x.out_dim)
+    return _kernelmatrix_kroneckerjl_helper(x, Kfeatures, Koutputs)
+end
+
+function kronecker_kernelmatrix(
+    k::MOKernel, x::IsotopicMOInputsUnion, y::IsotopicMOInputsUnion
+)
+    return throw(
+        ArgumentError("This kernel does not support a lazy kronecker kernelmatrix.")
+    )
+end
+
+function kronecker_kernelmatrix(k::MOKernel, x::IsotopicMOInputsUnion)
+    return kronecker_kernelmatrix(k, x, x)
+end

src/mokernels/independent.jl

@@ -27,41 +27,28 @@ function (κ::IndependentMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{Any,I
     return κ.kernel(x, y) * (px == py)
 end
 
-function _kernelmatrix_kron_helper(::MOInputIsotopicByFeatures, Kfeatures, B)
-    return kron(Kfeatures, B)
-end
-
-function _kernelmatrix_kron_helper(::MOInputIsotopicByOutputs, Kfeatures, B)
-    return kron(B, Kfeatures)
-end
+_mo_output_covariance(k::IndependentMOKernel, out_dim) = Eye{Bool}(out_dim)
 
 function kernelmatrix(
-    k::IndependentMOKernel, x::MOI, y::MOI
-) where {MOI<:IsotopicMOInputsUnion}
+    k::IndependentMOKernel, x::IsotopicMOInputsUnion, y::IsotopicMOInputsUnion
+)
     @assert x.out_dim == y.out_dim
     Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
-    mtype = eltype(Kfeatures)
-    return _kernelmatrix_kron_helper(x, Kfeatures, Eye{mtype}(x.out_dim))
+    Koutputs = _mo_output_covariance(k, x.out_dim)
+    return _kernelmatrix_kron_helper(x, Kfeatures, Koutputs)
 end
 
 if VERSION >= v"1.6"
-    function _kernelmatrix_kron_helper!(K, ::MOInputIsotopicByFeatures, Kfeatures, B)
-        return kron!(K, Kfeatures, B)
-    end
-
-    function _kernelmatrix_kron_helper!(K, ::MOInputIsotopicByOutputs, Kfeatures, B)
-        return kron!(K, B, Kfeatures)
-    end
-
     function kernelmatrix!(
-        K::AbstractMatrix, k::IndependentMOKernel, x::MOI, y::MOI
-    ) where {MOI<:IsotopicMOInputsUnion}
+        K::AbstractMatrix,
+        k::IndependentMOKernel,
+        x::IsotopicMOInputsUnion,
+        y::IsotopicMOInputsUnion,
+    )
         @assert x.out_dim == y.out_dim
-        Ktmp = kernelmatrix(k.kernel, x.x, y.x)
-        mtype = eltype(Ktmp)
-        return _kernelmatrix_kron_helper!(
-            K, x, Ktmp, Matrix{mtype}(I, x.out_dim, x.out_dim)
-        )
+        Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
+        Koutputs = _mo_output_covariance(k, x.out_dim)
+        return _kernelmatrix_kron_helper!(K, x, Kfeatures, Koutputs)
     end
 end
 

src/mokernels/intrinsiccoregion.jl

@@ -42,21 +42,31 @@ function (k::IntrinsicCoregionMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{
     return k.B[px, py] * k.kernel(x, y)
 end
 
+function _mo_output_covariance(k::IntrinsicCoregionMOKernel, out_dim)
+    @assert size(k.B) == (out_dim, out_dim)
+    return k.B
+end
+
 function kernelmatrix(
-    k::IntrinsicCoregionMOKernel, x::MOI, y::MOI
-) where {MOI<:IsotopicMOInputsUnion}
+    k::IntrinsicCoregionMOKernel, x::IsotopicMOInputsUnion, y::IsotopicMOInputsUnion
+)
     @assert x.out_dim == y.out_dim
     Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
-    return _kernelmatrix_kron_helper(x, Kfeatures, k.B)
+    Koutputs = _mo_output_covariance(k, x.out_dim)
+    return _kernelmatrix_kron_helper(x, Kfeatures, Koutputs)
 end
 
 if VERSION >= v"1.6"
     function kernelmatrix!(
-        K::AbstractMatrix, k::IntrinsicCoregionMOKernel, x::MOI, y::MOI
-    ) where {MOI<:IsotopicMOInputsUnion}
+        K::AbstractMatrix,
+        k::IntrinsicCoregionMOKernel,
+        x::IsotopicMOInputsUnion,
+        y::IsotopicMOInputsUnion,
+    )
         @assert x.out_dim == y.out_dim
         Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
-        return _kernelmatrix_kron_helper!(K, x, Kfeatures, k.B)
+        Koutputs = _mo_output_covariance(k, x.out_dim)
+        return _kernelmatrix_kron_helper!(K, x, Kfeatures, Koutputs)
     end
 end
 

src/mokernels/mokernel.jl

@@ -4,3 +4,21 @@
 Abstract type for kernels with multiple outpus.
 """
 abstract type MOKernel <: Kernel end
+
+function _kernelmatrix_kron_helper(::MOInputIsotopicByFeatures, Kfeatures, Koutputs)
+    return kron(Kfeatures, Koutputs)
+end
+
+function _kernelmatrix_kron_helper(::MOInputIsotopicByOutputs, Kfeatures, Koutputs)
+    return kron(Koutputs, Kfeatures)
+end
+
+if VERSION >= v"1.6"
+    function _kernelmatrix_kron_helper!(K, ::MOInputIsotopicByFeatures, Kfeatures, Koutputs)
+        return kron!(K, Kfeatures, Koutputs)
+    end
+
+    function _kernelmatrix_kron_helper!(K, ::MOInputIsotopicByOutputs, Kfeatures, Koutputs)
+        return kron!(K, Koutputs, Kfeatures)
+    end
+end

test/matrix/kernelkroneckermat.jl

@@ -7,4 +7,50 @@
     @test all(collect(kernelkronmat(k, collect(x), 2)) .≈ kernelmatrix(k, X; obsdim=1))
     @test all(collect(kernelkronmat(k, [x, x])) .≈ kernelmatrix(k, X; obsdim=1))
     @test_throws AssertionError kernelkronmat(LinearKernel(), collect(x), 2)
+
+    @testset "lazy kernelmatrix" begin
+        rng = MersenneTwister(123)
+
+        dims = (in=3, out=2, obs=3)
+        r = 1
+
+        A = randn(dims.out, r)
+        B = A * transpose(A) + Diagonal(rand(dims.out))
+
+        # XIF = [(rand(dims.in), rand(1:(dims.out))) for i in 1:(dims.obs)]
+        x = [rand(dims.in) for _ in 1:2]
+        XIF = KernelFunctions.MOInputIsotopicByFeatures(x, dims.out)
+        XIO = KernelFunctions.MOInputIsotopicByOutputs(x, dims.out)
+        y = [rand(dims.in) for _ in 1:2]
+        YIF = KernelFunctions.MOInputIsotopicByFeatures(y, dims.out)
+        YIO = KernelFunctions.MOInputIsotopicByOutputs(y, dims.out)
+
+        skernel = GaussianKernel()
+        kIndMO = IndependentMOKernel(skernel)
+
+        A = randn(dims.out, r)
+        B = A * transpose(A) + Diagonal(rand(dims.out))
+        icoregionkernel = IntrinsicCoregionMOKernel(skernel, B)
+
+        function test_kronecker_kernelmatrix(k, x)
+            res = kronecker_kernelmatrix(k, x)
+            @test typeof(res) <: Kronecker.KroneckerProduct
+            @test res == kernelmatrix(k, x)
+        end
+        function test_kronecker_kernelmatrix(k, x, y)
+            res = kronecker_kernelmatrix(k, x, y)
+            @test typeof(res) <: Kronecker.KroneckerProduct
+            @test res == kernelmatrix(k, x, y)
+        end
+
+        for k in [kIndMO, icoregionkernel], x in [XIF, XIO]
+            test_kronecker_kernelmatrix(k, x)
+        end
+        for k in [kIndMO, icoregionkernel], (x, y) in ([XIF, YIF], [XIO, YIO])
+            test_kronecker_kernelmatrix(k, x, y)
+        end
+
+        struct TestMOKernel <: MOKernel end
+        @test_throws ArgumentError kronecker_kernelmatrix(TestMOKernel(), XIF)
+    end
 end