Resolve type stability of evaluation of KernelSum (JuliaGaussianProcesses#459)

* Test type stability of evaluation of `KernelSum`

* Add hack

* Implement @willtebbutt's suggestion

* Revert strange hack

* Implement reviewdog's suggestion

* Apply reviewdog's spacing preferences

* Add comment

Co-authored-by: Will Tebbutt <[email protected]>

* Fix tests

* Bump patch

Co-authored-by: Will Tebbutt <[email protected]>
Co-authored-by: Will Tebbutt <[email protected]>
Co-authored-by: Will Tebbutt <[email protected]>

Author: 3 people

Project.toml

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

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,24 @@
     end
 
     test_params(k1 + k2, (k1, k2))
+
+    # Regression tests for https://github.com//issues/458
+    @testset "Type stability" begin
+        function check_type_stability(k)
+            @test (@inferred k(0.1, 0.2)) isa Real
+            x = rand(10)
+            y = rand(10)
+            @test (@inferred kernelmatrix(k, x)) isa Matrix{<:Real}
+            @test (@inferred kernelmatrix(k, x, y)) isa Matrix{<:Real}
+            @test (@inferred kernelmatrix_diag(k, x)) isa Vector{<:Real}
+            @test (@inferred kernelmatrix_diag(k, x, y)) isa Vector{<:Real}
+        end
+        @testset for k in (
+            RBFKernel() + RBFKernel() * LinearKernel(),
+            RBFKernel() + RBFKernel() * ExponentialKernel(),
+            RBFKernel() * (LinearKernel() + ExponentialKernel()),
+        )
+            check_type_stability(k)
+        end
+    end
 end