Repository: cornellius-gp/gpytorch
Issue: [Bug] Batch-indexing lazy covariance matrix fails for batch-mode ExactGP

# 🐛 Bug

Evaluating a batch-mode ExactGP with training inputs of shape `b x n x d` and targets of shape `b x n` on a set of test points with shape `m x d` yields a lazy_covariance_matrix with shape `b x m x m`. Indexing along the batch dimension fails.

## To reproduce

[batch_indexing_bug_lazy_covar.ipynb.txt](https://github.com/cornellius-gp/gpytorch/files/3331406/batch_indexing_bug_lazy_covar.ipynb.txt)

## Expected Behavior

Selecting a single index of the batch dimension:

```
mvn.lazy_covariance_matrix[0, :, :]
```
should yield the lazy covariance matrix for the 0th batch (with shape `m x m`)

## System information
- <!-- GPyTorch Version (run `print(gpytorch.__version__)` --> 0.3.2
- <!-- PyTorch Version (run `print(torch.__version__)` --> 1.0.0a0


Changed files:
test/kernels/test_scale_kernel.py, gpytorch/kernels/lcm_kernel.py, gpytorch/test/base_kernel_test_case.py, gpytorch/lazy/lazy_evaluated_kernel_tensor.py, gpytorch/kernels/kernel.py

diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
index 11414ca4..deb74fff 100644
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -10,6 +10,7 @@ from ..models.exact_prediction_strategies import DefaultPredictionStrategy, SumP
 from .. import settings
 from ..utils.deprecation import _ClassWithDeprecatedBatchSize, _deprecate_kwarg_with_transform
 from ..constraints import Positive
+from copy import deepcopy
 
 
 def default_postprocess_script(x):
@@ -383,6 +384,26 @@ class Kernel(Module, _ClassWithDeprecatedBatchSize):
     def __setstate__(self, d):
         self.__dict__ = d
 
+    def __getitem__(self, index):
+        if len(self.batch_shape) == 0:
+            return self
+
+        new_kernel = deepcopy(self)
+        # Process the index
+        index = index if isinstance(index, tuple) else (index,)
+
+        for param_name, param in self._parameters.items():
+            new_kernel._parameters[param_name].data = param.__getitem__(index)
+            ndim_removed = len(param.shape) - len(new_kernel._parameters[param_name].shape)
+            new_batch_shape_len = len(self.batch_shape) - ndim_removed
+            new_kernel.batch_shape = new_kernel._parameters[param_name].shape[:new_batch_shape_len]
+
+        for sub_module_name, sub_module in self._modules.items():
+            if isinstance(sub_module, Kernel):
+                self._modules[sub_module_name] = sub_module.__getitem__(index)
+
+        return new_kernel
+
 
 class AdditiveKernel(Kernel):
     """
@@ -415,6 +436,13 @@ class AdditiveKernel(Kernel):
     def num_outputs_per_input(self, x1, x2):
         return self.kernels[0].num_outputs_per_input(x1, x2)
 
+    def __getitem__(self, index):
+        new_kernel = deepcopy(self)
+        for i, kernel in enumerate(self.kernels):
+            new_kernel.kernels[i] = self.kernels[i].__getitem__(index)
+
+        return new_kernel
+
 
 class ProductKernel(Kernel):
     """
@@ -457,3 +485,10 @@ class ProductKernel(Kernel):
 
     def num_outputs_per_input(self, x1, x2):
         return self.kernels[0].num_outputs_per_input(x1, x2)
+
+    def __getitem__(self, index):
+        new_kernel = deepcopy(self)
+        for i, kernel in enumerate(self.kernels):
+            new_kernel.kernels[i] = self.kernels[i].__getitem__(index)
+
+        return new_kernel
diff --git a/gpytorch/kernels/lcm_kernel.py b/gpytorch/kernels/lcm_kernel.py
index d10d38c9..783231fb 100644
--- a/gpytorch/kernels/lcm_kernel.py
+++ b/gpytorch/kernels/lcm_kernel.py
@@ -3,6 +3,7 @@
 from torch.nn import ModuleList
 from .kernel import Kernel
 from .multitask_kernel import MultitaskKernel
+from copy import deepcopy
 
 
 class LCMKernel(Kernel):
@@ -50,3 +51,10 @@ class LCMKernel(Kernel):
         returns an `(n*num_tasks) x (m*num_tasks)` covariance matrix.
         """
         return self.covar_module_list[0].num_outputs_per_input(x1, x2)
+
+    def __getitem__(self, index):
+        new_kernel = deepcopy(self)
+        new_kernel.covar_module_list = ModuleList(
+            [base_kernel.__getitem__(index) for base_kernel in self.covar_module_list]
+        )
+        return new_kernel
diff --git a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
index c9721b2e..ce4226fa 100644
--- a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
+++ b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
@@ -117,8 +117,13 @@ class LazyEvaluatedKernelTensor(LazyTensor):
                 x2 = x2.expand(*([1] * len(batch_indices)), *self.x2.shape[-2:])
                 x2 = x2[(*batch_indices, row_index, dim_index)]
 
+        if len(batch_indices) == 0 or all(ind == slice(None, None, None) for ind in batch_indices):
+            new_kernel = self.kernel  # Avoid unnecessary copying when we aren't explicitly indexing batch dims
+        else:
+            new_kernel = self.kernel.__getitem__(batch_indices)
+
         # Now construct a kernel with those indices
-        return self.__class__(x1, x2, kernel=self.kernel, last_dim_is_batch=self.last_dim_is_batch, **self.params)
+        return self.__class__(x1, x2, kernel=new_kernel, last_dim_is_batch=self.last_dim_is_batch, **self.params)
 
     def _matmul(self, rhs):
         # This _matmul is defined computes the kernel in chunks
diff --git a/gpytorch/test/base_kernel_test_case.py b/gpytorch/test/base_kernel_test_case.py
index 77d2ea1b..1d6f40e7 100644
--- a/gpytorch/test/base_kernel_test_case.py
+++ b/gpytorch/test/base_kernel_test_case.py
@@ -114,3 +114,25 @@ class BaseKernelTestCase(object):
         x = self.create_data_double_batch()
         batch_covar_mat = kernel(x).evaluate_kernel().evaluate()
         return batch_covar_mat
+
+    def test_kernel_getitem_single_batch(self):
+        kernel = self.create_kernel_no_ard(batch_shape=torch.Size([2]))
+        x = self.create_data_single_batch()
+
+        res1 = kernel(x).evaluate()[0]  # Result of first kernel on first batch of data
+
+        new_kernel = kernel[0]
+        res2 = new_kernel(x[0]).evaluate()  # Should also be result of first kernel on first batch of data.
+
+        self.assertLess(torch.norm(res1 - res2) / res1.norm(), 1e-4)
+
+    def test_kernel_getitem_double_batch(self):
+        kernel = self.create_kernel_no_ard(batch_shape=torch.Size([3, 2]))
+        x = self.create_data_double_batch()
+
+        res1 = kernel(x).evaluate()[0, 1]  # Result of first kernel on first batch of data
+
+        new_kernel = kernel[0, 1]
+        res2 = new_kernel(x[0, 1]).evaluate()  # Should also be result of first kernel on first batch of data.
+
+        self.assertLess(torch.norm(res1 - res2) / res1.norm(), 1e-4)
diff --git a/test/kernels/test_scale_kernel.py b/test/kernels/test_scale_kernel.py
index 610a957f..e910df25 100644
--- a/test/kernels/test_scale_kernel.py
+++ b/test/kernels/test_scale_kernel.py
@@ -3,9 +3,20 @@
 import torch
 import unittest
 from gpytorch.kernels import RBFKernel, ScaleKernel
+from gpytorch.test.base_kernel_test_case import BaseKernelTestCase
 
 
-class TestScaleKernel(unittest.TestCase):
+class TestScaleKernel(BaseKernelTestCase, unittest.TestCase):
+    def create_kernel_no_ard(self, **kwargs):
+        base_kernel = RBFKernel()
+        kernel = ScaleKernel(base_kernel, **kwargs)
+        return kernel
+
+    def create_kernel_ard(self, num_dims, **kwargs):
+        base_kernel = RBFKernel(ard_num_dims=num_dims)
+        kernel = ScaleKernel(base_kernel, **kwargs)
+        return kernel
+
     def test_ard(self):
         a = torch.tensor([[1, 2], [2, 4]], dtype=torch.float)
         b = torch.tensor([[1, 3], [0, 4]], dtype=torch.float)