Type stability of KernelSum

[simsurace]

using KernelFunctions
using Test
k = RBFKernel() + RBFKernel() * ExponentialKernel()
@inferred k(0.1, 0.2) # ERROR
@time k(0.1, 0.2)  # 0.000009 seconds (5 allocations: 176 bytes)

but e.g.

using KernelFunctions
using Test
k = RBFKernel() + RBFKernel() * LinearKernel()
@inferred k(0.1, 0.2) # works

Similar results hold for kernelmatrix and kernelmatrix_diag.
I haven't figured out why certain kernels show the problem while others don't.
This may be a variant of JuliaLang/julia#45748 and could be a Julia issue.
The similarity to the referenced Julia issue is that calling sum before affects type inference:

# restart Julia
using KernelFunctions
using Test
k = RBFKernel() + RBFKernel() * ExponentialKernel()
sum(f -> f(0.1, 0.2), k.kernels)
@inferred k(0.1, 0.2) # works
@time k(0.1, 0.2) # 0.000008 seconds (1 allocation: 16 bytes)

Despite this probably being a Julia issue, we may want to patch it anyways, as it can adversely impact AD for certain GP models.
I'll open a PR.

[devmotion]

Despite this probably being a Julia issue, we may want to patch it anyways

What's a possible workaround? Can we just change

KernelFunctions.jl/src/kernels/kernelsum.jl

Line 46 in ce7923f

(κ::KernelSum)(x, y) = sum(k(x, y) for k in κ.kernels)

to

function (κ::KernelSum)(x, y)
    return sum(κ.kernels) do k
        return k(x, y)
    end
end

etc.?

[simsurace]

I'll double-check, but I think I tried that and it did not solve the issue. I think I'll use the init keyword, as the Julia issue seems to be related to the inability to figure that out automatically.

[devmotion]

Maybe one has to add a let block to avoid the closure type inference bug?

[simsurace]

Hmm, seemed to be a false alarm. The issue is that many things seem to work if the function is redefined after module import, but not in a new Julia session immediately after precompilation.

[willtebbutt]

@simsurace can we close this now?

[simsurace]

Sure. It should have been linked to #459 I guess, but it doesn't get linked automatically when non-admins use the closing keywords.