Repository: cornellius-gp/gpytorch

Issue:
[Bug] Lazy Tensor Internal Indexing Error for Large Variance Evaluation (MultitaskMultivariateNormal)
# 🐛 Bug
When calculating the `variance` of a `MultitaskMultivariateNormal`, evaluating `self.lazy_covariance_matrix.diag()` returns an indexing error when the number of query points is too large. When the number of data points is small (<150), inside the `MultitaskMultivariateNormal` we have
```
self.lazy_covariance_matrix.lazy_tensors
Out[2]: 
(<gpytorch.lazy.non_lazy_tensor.NonLazyTensor at 0x125d53350>,
 <gpytorch.lazy.added_diag_lazy_tensor.AddedDiagLazyTensor at 0x12cd99050>)
```
But when we increase the number of query points we get
```
self.lazy_covariance_matrix.lazy_tensors
Out[4]: 
(<gpytorch.lazy.interpolated_lazy_tensor.InterpolatedLazyTensor at 0x131338410>,
 <gpytorch.lazy.matmul_lazy_tensor.MatmulLazyTensor at 0x131338550>,
 <gpytorch.lazy.block_diag_lazy_tensor.BlockDiagLazyTensor at 0x131338290>,
 <gpytorch.lazy.diag_lazy_tensor.DiagLazyTensor at 0x131338610>)
```
Having different implementation strategy depending on size is fine, but the strategy for large batches errors out when evaluating variance.
## To reproduce
Pretty much the Batch Independent Multioutput GP [tutorial](https://gpytorch.readthedocs.io/en/latest/examples/03_Multitask_Exact_GPs/Batch_Independent_Multioutput_GP.html) but increase the `ntest` to 200 (no problems if you use 51 as in the tutorial)

** Code snippet to reproduce **
```python
import math
import torch
import gpytorch
from matplotlib import pyplot as plt

train_x = torch.linspace(0, 1, 100)
ntest = 200

train_y = torch.stack([
    torch.sin(train_x * (2 * math.pi)) + torch.randn(train_x.shape[0]) * 0.2,
    torch.cos(train_x * (2 * math.pi)) + torch.randn(train_x.shape[0]) * 0.8,
], -1)
ny = train_y.shape[1]


class BatchIndependentMultitaskGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super().__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([ny]))
        self.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.RBFKernel(batch_shape=torch.Size([ny])),
            batch_shape=torch.Size([ny])
        )

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultitaskMultivariateNormal.from_batch_mvn(
            gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
        )


likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=ny)
model = BatchIndependentMultitaskGPModel(train_x, train_y, likelihood)

training_iterations = 50

# Find optimal model hyperparameters
model.train()
likelihood.train()

# Use the adam optimizer
optimizer = torch.optim.Adam([
    {'params': model.parameters()},  # Includes GaussianLikelihood parameters
], lr=0.1)

# "Loss" for GPs - the marginal log likelihood
mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

for i in range(training_iterations):
    optimizer.zero_grad()
    output = model(train_x)
    loss = -mll(output, train_y)
    loss.backward()
    print('Iter %d/%d - Loss: %.3f' % (i + 1, training_iterations, loss.item()))
    optimizer.step()
# Set into eval mode
model.eval()
likelihood.eval()

# Initialize plots
f, axes = plt.subplots(1, ny, figsize=(8, 3))

# Make predictions
with torch.no_grad():
    test_x = torch.linspace(0, 1, ntest)
    xx = torch.stack([test_x, torch.linspace(0, 2, test_x.shape[0])], -1)
    predictions = likelihood(model(test_x))
    mean = predictions.mean
    lower, upper = predictions.confidence_region()

# This contains predictions for both tasks, flattened out
# The first half of the predictions is for the first task
# The second half is for the second task

for i in range(train_y.shape[1]):
    # Plot training data as black stars
    axes[i].plot(train_x.detach().numpy(), train_y[:, i].detach().numpy(), 'k*')
    # Predictive mean as blue line
    axes[i].plot(test_x.numpy(), mean[:, i].numpy(), 'b')
    # Shade in confidence
    axes[i].fill_between(test_x.numpy(), lower[:, i].numpy(), upper[:, i].numpy(), alpha=0.5)
    axes[i].set_ylim([-3, 3])
axes[-1].legend(['Observed Data', 'Mean', 'Confidence'])
axes[-1].set_title('Observed Values (Likelihood)')
plt.show()
```

** Stack trace/error message **
```
Traceback (most recent call last):
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py", line 88, in _getitem
    x1 = x1[(*batch_indices, row_index, dim_index)]
IndexError: too many indices for tensor of dimension 2

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/johnsonzhong/Research/meta_contact/simulation/temp.py", line 70, in <module>
    lower, upper = predictions.confidence_region()
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/distributions/multivariate_normal.py", line 80, in confidence_region
    std2 = self.stddev.mul_(2)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/torch/distributions/distribution.py", line 111, in stddev
    return self.variance.sqrt()
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/distributions/multitask_multivariate_normal.py", line 219, in variance
    var = super().variance
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/distributions/multivariate_normal.py", line 189, in variance
    diag = self.lazy_covariance_matrix.diag()
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/sum_lazy_tensor.py", line 96, in diag
    return sum(lazy_tensor.diag().contiguous() for lazy_tensor in self.lazy_tensors)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/sum_lazy_tensor.py", line 96, in <genexpr>
    return sum(lazy_tensor.diag().contiguous() for lazy_tensor in self.lazy_tensors)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/interpolated_lazy_tensor.py", line 387, in diag
    return super(InterpolatedLazyTensor, self).diag()
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py", line 839, in diag
    return self[..., row_col_iter, row_col_iter]
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py", line 1716, in __getitem__
    res = self._get_indices(row_index, col_index, *batch_indices)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/interpolated_lazy_tensor.py", line 110, in _get_indices
    *[batch_index.view(*batch_index.shape, 1, 1) for batch_index in batch_indices],
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/block_interleaved_lazy_tensor.py", line 58, in _get_indices
    res = self.base_lazy_tensor._get_indices(row_index, col_index, *batch_indices, row_index_block)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py", line 294, in _get_indices
    base_lazy_tensor = self._getitem(_noop_index, _noop_index, *batch_indices)._expand_batch(final_shape)
  File "/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py", line 97, in _getitem
    f"Attempting to tensor index a non-batch matrix's batch dimensions. "
RuntimeError: Attempting to tensor index a non-batch matrix's batch dimensions. Got batch index {batch_indices} but my shape was {self.shape}
```

## Expected Behavior
Calculating variance should work the same regardless of number of query points.

## System information

- <!-- GPyTorch Version (run `print(gpytorch.__version__)` --> gpytorch 1.01
- <!-- PyTorch Version (run `print(torch.__version__)` --> pytorch 1.3.1
- <!-- Computer OS --> MacOS 10.15.3 




Discussion:
@jacobrgardner, @gpleiss for 100 train points this happens exactly when switching eval points form 156 -> 157. I.e. exactly when the full covariance matrix grows > 256. Probably not a coincidence... 
Yup I suspect that it's switching implementation for larger matrices (suggested by the change in number of self.lazy_tensors), which is totally fine, but there's something wrong with the large matrix implementation
This seems to be only an issue with the BatchIndependentMultitaskGPModel from the tutorial. The Multitask one (allow correlated outputs) works fine.
:( My guess is this is probably related to #1024 and #1039 ... I'll work on this issue tomorrow.

Resolving patch:
diff --git a/gpytorch/lazy/block_lazy_tensor.py b/gpytorch/lazy/block_lazy_tensor.py
--- a/gpytorch/lazy/block_lazy_tensor.py
+++ b/gpytorch/lazy/block_lazy_tensor.py
@@ -4,6 +4,7 @@
 
 import torch
 
+from ..utils.getitem import _is_noop_index, _noop_index
 from .lazy_tensor import LazyTensor
 from .non_lazy_tensor import lazify
 
@@ -60,6 +61,42 @@ def _expand_batch(self, batch_shape):
         res = self.__class__(self.base_lazy_tensor._expand_batch(batch_shape))
         return res
 
+    def _getitem(self, row_index, col_index, *batch_indices):
+        # First the easy case: just batch indexing
+        if _is_noop_index(row_index) and _is_noop_index(col_index):
+            return self.__class__(self.base_lazy_tensor._getitem(row_index, col_index, *batch_indices, _noop_index))
+
+        # If either of the dimensions are indices, it's too complicated - go with the base case
+        if not isinstance(row_index, slice) or not isinstance(col_index, slice):
+            # It's too complicated to deal with tensor indices in this case - we'll use the super method
+            return super()._getitem(row_index, col_index, *batch_indices)
+
+        # Now we know that row_index and col_index
+        num_blocks = self.num_blocks
+        num_rows, num_cols = self.matrix_shape
+        row_start, row_end, row_step = row_index.start or 0, row_index.stop or num_rows, row_index.step
+        col_start, col_end, col_step = col_index.start or 0, col_index.stop or num_cols, col_index.step
+
+        # If we have a step, it's too complicated - go with the base case
+        if row_step is not None or col_step is not None:
+            return super()._getitem(row_index, col_index, *batch_indices)
+
+        # Let's make sure that the slice dimensions perfectly correspond with the number of
+        # outputs per input that we have
+        # Otherwise - its too complicated. We'll go with the base case
+        if (row_start % num_blocks) or (col_start % num_blocks) or (row_end % num_blocks) or (col_end % num_blocks):
+            return super()._getitem(row_index, col_index, *batch_indices)
+
+        # Otherwise - let's divide the slices by the number of outputs per input
+        row_index = slice(row_start // num_blocks, row_end // num_blocks, None)
+        col_index = slice(col_start // num_blocks, col_end // num_blocks, None)
+
+        # Now we can try the super call!
+        new_base_lazy_tensor = self.base_lazy_tensor._getitem(row_index, col_index, *batch_indices)
+
+        # Now construct a kernel with those indices
+        return self.__class__(new_base_lazy_tensor, block_dim=-3)
+
     def _matmul(self, rhs):
         isvector = rhs.ndimension() == 1
         if isvector:
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -17,7 +17,7 @@
 from ..utils.broadcasting import _matmul_broadcast_shape, _mul_broadcast_shape
 from ..utils.cholesky import psd_safe_cholesky
 from ..utils.deprecation import _deprecate_renamed_methods
-from ..utils.getitem import _compute_getitem_size, _convert_indices_to_tensors, _noop_index
+from ..utils.getitem import _compute_getitem_size, _convert_indices_to_tensors, _is_noop_index, _noop_index
 from ..utils.gradients import _ensure_symmetric_grad
 from ..utils.memoize import add_to_cache, cached
 from .lazy_tensor_representation_tree import LazyTensorRepresentationTree
@@ -213,7 +213,7 @@ def _getitem(self, row_index, col_index, *batch_indices):
             `LazyTensor`
         """
         # Special case: if both row and col are not indexed, then we are done
-        if row_index is _noop_index and col_index is _noop_index:
+        if _is_noop_index(row_index) and _is_noop_index(col_index):
             if len(batch_indices):
                 components = [component[batch_indices] for component in self._args]
                 res = self.__class__(*components, **self._kwargs)
diff --git a/gpytorch/utils/getitem.py b/gpytorch/utils/getitem.py
--- a/gpytorch/utils/getitem.py
+++ b/gpytorch/utils/getitem.py
@@ -168,6 +168,13 @@ def _equal_indices(a, b):
         return False
 
 
+def _is_noop_index(index):
+    """
+    Determine if a given index is a noop (e.g. ":")
+    """
+    return isinstance(index, slice) and index == _noop_index
+
+
 def _is_tensor_index_moved_to_start(indices):
     """
     Given an index, determine if the indexed part of the getitem is moved to the zero'th dimension


Test patch:
diff --git a/test/examples/test_independent_multitask_gp_regression.py b/test/examples/test_independent_multitask_gp_regression.py
new file mode 100644
--- /dev/null
+++ b/test/examples/test_independent_multitask_gp_regression.py
@@ -0,0 +1,104 @@
+#!/usr/bin/env python3
+
+import math
+import os
+import random
+import unittest
+
+import gpytorch
+import torch
+from gpytorch.distributions import MultivariateNormal, MultitaskMultivariateNormal
+from gpytorch.kernels import RBFKernel, ScaleKernel
+from gpytorch.likelihoods import MultitaskGaussianLikelihood
+from gpytorch.means import ConstantMean
+from torch import optim
+
+
+# Batch training test: Let's learn hyperparameters on a sine dataset, but test on a sine dataset and a cosine dataset
+# in parallel.
+train_x = torch.linspace(0, 2, 51).unsqueeze(-1)
+test_x = torch.linspace(0, 2, 11).unsqueeze(-1)
+train_y1 = torch.sin(train_x * (2 * math.pi)).squeeze()
+train_y1.add_(torch.randn_like(train_y1).mul_(0.01))
+test_y1 = torch.sin(test_x * (2 * math.pi)).squeeze()
+train_y2 = torch.sin(train_x * (2 * math.pi)).squeeze()
+train_y2.add_(torch.randn_like(train_y2).mul_(0.01))
+test_y2 = torch.sin(test_x * (2 * math.pi)).squeeze()
+
+# Combined sets of data
+train_y12 = torch.stack((train_y1, train_y2), dim=-1).contiguous()
+test_y12 = torch.stack((test_y1, test_y2), dim=-1).contiguous()
+
+
+class ExactGPModel(gpytorch.models.ExactGP):
+    def __init__(self, train_inputs, train_targets, likelihood, batch_shape=torch.Size([2])):
+        super(ExactGPModel, self).__init__(train_inputs, train_targets, likelihood)
+        self.mean_module = ConstantMean(batch_shape=batch_shape)
+        self.covar_module = ScaleKernel(RBFKernel(batch_shape=batch_shape), batch_shape=batch_shape)
+
+    def forward(self, x):
+        mean_x = self.mean_module(x)
+        covar_x = self.covar_module(x)
+        return MultitaskMultivariateNormal.from_batch_mvn(MultivariateNormal(mean_x, covar_x))
+
+
+class TestIndependentMultitaskGPRegression(unittest.TestCase):
+    def setUp(self):
+        if os.getenv("UNLOCK_SEED") is None or os.getenv("UNLOCK_SEED").lower() == "false":
+            self.rng_state = torch.get_rng_state()
+            torch.manual_seed(0)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed_all(0)
+            random.seed(0)
+
+    def tearDown(self):
+        if hasattr(self, "rng_state"):
+            torch.set_rng_state(self.rng_state)
+
+    def test_train_and_eval(self):
+        # We're manually going to set the hyperparameters to something they shouldn't be
+        likelihood = MultitaskGaussianLikelihood(num_tasks=2)
+        gp_model = ExactGPModel(train_x, train_y12, likelihood)
+        mll = gpytorch.ExactMarginalLogLikelihood(likelihood, gp_model)
+
+        # Find optimal model hyperparameters
+        gp_model.train()
+        likelihood.train()
+        optimizer = optim.Adam(gp_model.parameters(), lr=0.1)
+        optimizer.n_iter = 0
+        for _ in range(75):
+            optimizer.zero_grad()
+            output = gp_model(train_x)
+            loss = -mll(output, train_y12).sum()
+            loss.backward()
+            optimizer.step()
+
+            for param in gp_model.parameters():
+                self.assertTrue(param.grad is not None)
+                self.assertGreater(param.grad.norm().item(), 0)
+            for param in likelihood.parameters():
+                self.assertTrue(param.grad is not None)
+                self.assertGreater(param.grad.norm().item(), 0)
+
+        # Test the model
+        gp_model.eval()
+        likelihood.eval()
+
+        # Make predictions for both sets of test points, and check MAEs.
+        with torch.no_grad(), gpytorch.settings.max_eager_kernel_size(1):
+            batch_predictions = likelihood(gp_model(test_x))
+            preds1 = batch_predictions.mean[:, 0]
+            preds2 = batch_predictions.mean[:, 1]
+            mean_abs_error1 = torch.mean(torch.abs(test_y1 - preds1))
+            mean_abs_error2 = torch.mean(torch.abs(test_y2 - preds2))
+            self.assertLess(mean_abs_error1.squeeze().item(), 0.01)
+            self.assertLess(mean_abs_error2.squeeze().item(), 0.01)
+
+            # Smoke test for getting predictive uncertainties
+            lower, upper = batch_predictions.confidence_region()
+            self.assertEqual(lower.shape, test_y12.shape)
+            self.assertEqual(upper.shape, test_y12.shape)
+
+
+if __name__ == "__main__":
+    unittest.main()