[Feature Request] Batched GridInterpolationKernel

[Ilykuleshov]

🚀 Feature Request

If I'm not missing anything, there is currently no way to use the GridInterpolationKernel with a batched input. The only alternative currently (correct me if I'm wrong) is to iterate over the batch dimension, calculating the kernel individually for each element, which is suboptimal. Would be great to see a batched GridInterpolationKernel, seems like a very useful feature.

Motivation

You have already implemented batch shape for the RBF Kernel, its only natural to extend it to the rest of the kernels.

Pitch

This would greatly increase useability. AFAIK this is the only library that implements SKI/KISS-GP, so this would greatly improve useability.

I've tried to wrap it in vmap, but it fails due to some .item call.