KernelFunctions.jl icon indicating copy to clipboard operation
KernelFunctions.jl copied to clipboard

Type stability of `KernelSum`

Open simsurace opened this issue 2 years ago • 4 comments

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 https://github.com/JuliaLang/julia/issues/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.

simsurace avatar Jun 23 '22 11:06 simsurace

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

What's a possible workaround? Can we just change https://github.com/JuliaGaussianProcesses/KernelFunctions.jl/blob/ce7923f8252d183768f247256079a2c8279f7b1b/src/kernels/kernelsum.jl#L46 to

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

etc.?

devmotion avatar Jun 23 '22 12:06 devmotion

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.

simsurace avatar Jun 23 '22 12:06 simsurace

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

devmotion avatar Jun 23 '22 13:06 devmotion

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.

simsurace avatar Jun 23 '22 14:06 simsurace

@simsurace can we close this now?

willtebbutt avatar Oct 15 '22 12:10 willtebbutt

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

simsurace avatar Oct 17 '22 16:10 simsurace