SelectTransform with single index

[willtebbutt]

The output SelectTransform(5)(some_vector) should really be a scalar, however, this currently doesn't work. It is important that this works so that (for example) a single dimension of a vector-valued input can be passed to kernels that can't handle vector-valued inputs, such as kernels constructed via PeriodicTransform.

This PR addresses this issue by making this possible.

[willtebbutt]

@molet could you check that this fixes your problem?

[molet]

@molet could you check that this fixes your problem?

Thanks a lot @willtebbutt for working on it! Now it seems to be working nicely:

using KernelFunctions, Zygote, AxisArrays
n = 100;
data = Matrix(hcat(collect(1:n), rand(n))');
a = AxisArray(data, features=[:time, :load], idx=collect(1:n));
k(θ) = TransformedKernel(TransformedKernel(
    TransformedKernel(ExponentialKernel(), ScaleTransform(θ[1])), 
    PeriodicTransform(1/24.0)),
    SelectTransform(:time)
);
o(θ) = sum(kernelmatrix(k(θ), a));
o([2.3])
1515.157184696975
Zygote.gradient(o, [2.3])
([294.83700581618575],)

[molet]

LGTM

[willtebbutt]

@theogf anything else you need?

[theogf]

Nope it's all good!