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



Discussion:
This issue comes from having a batch-mode kernel, which is not "indexed" along the batch dimension when the LazyEvaluatedKernelTensor is indexed along the batch dimension--for this use case, the kernel would also need to be "indexed" along the batch dim when creating the new LazyEvaluatedKernelTensor: https://github.com/cornellius-gp/gpytorch/blob/master/gpytorch/lazy/lazy_evaluated_kernel_tensor.py#L121
@jacobrgardner, @gpleiss is this something you think you can fix in the near future? This is kind of important in order to get multi-output models working efficiently.

Resolving patch:
diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -10,6 +10,7 @@
 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 @@ def __mul__(self, other):
     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 @@ def prediction_strategy(self, train_inputs, train_prior_dist, train_labels, like
     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 @@ def forward(self, x1, x2, diag=False, **params):
 
     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
--- 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 @@ def num_outputs_per_input(self, x1, x2):
         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
--- a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
+++ b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
@@ -117,8 +117,13 @@ def _getitem(self, row_index, col_index, *batch_indices):
                 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


Test patch:
diff --git a/gpytorch/test/base_kernel_test_case.py b/gpytorch/test/base_kernel_test_case.py
--- a/gpytorch/test/base_kernel_test_case.py
+++ b/gpytorch/test/base_kernel_test_case.py
@@ -114,3 +114,25 @@ def test_smoke_double_batch_kernel_double_batch_x_ard(self):
         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
--- 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)