Fix issues with indexing broadcasted LazyEvaluatedKernelTensors

[Balandat]

When batch-evaluating the posterior of a model (i.e.when, in eval mode, passing in a tensor with a batch shape that has more dimensions than the batch shape of the training data), then, if the kernel is lazily evaluated, this leads to obscure indexing errors. The issue is that the raw parameters of the lazy evaluated kernel have different batch shapes than the batch shape of the covariance matrix.

Specifically, if train_batch_shape is the batch shape of the kernel (unbroadcasted, as used for computing the train-train covariance)and the posterior is evaluated with some batch shape that has more dimensions, e.g.eval_batch_shape = add_batch_shape + train_batch_shape, then under the hood when the LazyEvaluatedKernelTensoris indexed, we use a batch index meant for theeval_batch_shapeto index the underlyig ernel parameters (e.g. lengthscales) who only have the leadingtrain_batch_shape, which casues indexing issues. This doesn't happen in the case of eager evaluations of the kernel, since the parameters are appropriately broadcasted to the new larger eval_batch_shapebefore indexing happens (i.e. when computing the kernel the parameters are automatically "broadcasted up" to the shape that's required to produce the kernel with theeval_batch_shape`.

It's not entirely clear how to best fix this. The path pursued with this PR is to check pre-indexing whether the dimension of the batch shapes are different, and then automatically modify the indices that are passed to the __getitem__ function of kernel (and thus subsequently to the various lazy tensors that the kernel consists of). It appears to work but seems quite dicey and easy to break, so open for alternative suggestions.

There likely will be similar issues if the posterior of a batchd model evaluated on a non-batch input is indexed while lazilyt evaluated, but that's something for an separate PR.

[Balandat]

@abbas-kazerouni

[dme65]

lgtm!