Resolve type stability of evaluation of KernelSum

src/kernels/kernelsum.jl

@@ -43,22 +43,25 @@ end
 
 Base.length(k::KernelSum) = length(k.kernels)
 
-(κ::KernelSum)(x, y) = sum(k(x, y) for k in κ.kernels)
+_sum(f::Tf, x::Tuple) where {Tf} = f(x[1]) + _sum(f, Base.tail(x))
+_sum(f::Tf, x::Tuple{Tx}) where {Tf, Tx} = f(x[1])
+
+(κ::KernelSum)(x, y) = _sum(k -> k(x, y), κ.kernels)
 
 function kernelmatrix(κ::KernelSum, x::AbstractVector)
-    return sum(kernelmatrix(k, x) for k in κ.kernels)
+    return _sum(Base.Fix2(kernelmatrix, x), κ.kernels)
 end
 
 function kernelmatrix(κ::KernelSum, x::AbstractVector, y::AbstractVector)
-    return sum(kernelmatrix(k, x, y) for k in κ.kernels)
+    return _sum(k -> kernelmatrix(k, x, y), κ.kernels)
 end
 
 function kernelmatrix_diag(κ::KernelSum, x::AbstractVector)
-    return sum(kernelmatrix_diag(k, x) for k in κ.kernels)
+    return _sum(Base.Fix2(kernelmatrix_diag, x), κ.kernels)
 end
 
 function kernelmatrix_diag(κ::KernelSum, x::AbstractVector, y::AbstractVector)
-    return sum(kernelmatrix_diag(k, x, y) for k in κ.kernels)
+    return _sum(k -> kernelmatrix_diag(k, x, y), κ.kernels)
 end
 
 function Base.show(io::IO, κ::KernelSum)

test/kernels/kernelsum.jl

@@ -18,4 +18,23 @@
     end
 
     test_params(k1 + k2, (k1, k2))
+
+    @testset "Type stability" begin
+        function check_type_stability(k)
+            @inferred k(0.1, 0.2)
+            x = rand(10)
+            y = rand(10)
+            @inferred kernelmatrix(k, x)
+            @inferred kernelmatrix(k, x, y)
+            @inferred kernelmatrix_diag(k, x)
+            @inferred kernelmatrix_diag(k, x, y)
+        end
+        @testset for k in (
+            RBFKernel() + RBFKernel() * LinearKernel(),
+            RBFKernel() + RBFKernel() * ExponentialKernel(),
+            RBFKernel() * (LinearKernel() + ExponentialKernel())
+        )
+            check_type_stability(k)
+        end
+    end
 end