Repository: cornellius-gp/gpytorch
Pull request #404 (2018-11-29 15:39:36)

Improve LazyTensor.__getitem__
A more stable version of yesterday's hotfix.

`LazyTensor.__getitem__` now calls a `LazyTensor._getitem` helper (similar to how `LazyTensor.size` calls `LazyTensor._size`). When `LazyTensor._getitem` is called, there is:
- One index per dimension (unsupplied indices are filled with `slice(None, None, None)`.
- Each index is an int, a slice, or a LongTensor (all lists, Ellipsis have been processed out).

LTs that had custom `__getitem__`s now have custom `_getitem`s. This is a pure refactor - no new functionality (except for `ConstantMulLazyTensor`, which is now a bit more stable). This significantly reduces repeated code.


Patch:
diff --git a/gpytorch/lazy/batch_repeat_lazy_tensor.py b/gpytorch/lazy/batch_repeat_lazy_tensor.py
index d2400fcd..950101cd 100644
--- a/gpytorch/lazy/batch_repeat_lazy_tensor.py
+++ b/gpytorch/lazy/batch_repeat_lazy_tensor.py
@@ -24,6 +24,23 @@ class BatchRepeatLazyTensor(LazyTensor):
         batch_indices = batch_indices[-num_true_batch_dims:] if num_true_batch_dims else []
         return self.base_lazy_tensor._get_indices(row_indices, col_indices, *batch_indices)
 
+    def _getitem(self, *indices):
+        args = []
+        kwargs = self.base_lazy_tensor._kwargs
+        num_base_batch_dims = len(self.base_lazy_tensor.batch_shape)
+
+        for arg in self.base_lazy_tensor._args:
+            if torch.is_tensor(arg):
+                arg_base_shape_len = max(arg.dim() - num_base_batch_dims, 0)
+                args.append(arg.repeat(*self.batch_repeat, *[1 for _ in range(arg_base_shape_len)]))
+            elif isinstance(arg, LazyTensor):
+                args.append(BatchRepeatLazyTensor(arg, batch_repeat=self.batch_repeat))
+            else:
+                args.append(arg)
+
+        new_lazy_tensor = self.base_lazy_tensor.__class__(*args, **kwargs)
+        return new_lazy_tensor._getitem(*indices)
+
     def _matmul(self, rhs):
         rhs = self._move_repeat_batches_to_columns(rhs)
         res = self.base_lazy_tensor._matmul(rhs)
@@ -167,20 +184,3 @@ class BatchRepeatLazyTensor(LazyTensor):
                 initial_vectors=initial_vectors, test_vectors=test_vectors
             ).root.repeat(*self.batch_repeat, 1, 1)
         )
-
-    def __getitem__(self, index):
-        args = []
-        kwargs = self.base_lazy_tensor._kwargs
-        num_base_batch_dims = len(self.base_lazy_tensor.batch_shape)
-
-        for arg in self.base_lazy_tensor._args:
-            if torch.is_tensor(arg):
-                arg_base_shape_len = max(arg.dim() - num_base_batch_dims, 0)
-                args.append(arg.repeat(*self.batch_repeat, *[1 for _ in range(arg_base_shape_len)]))
-            elif isinstance(arg, LazyTensor):
-                args.append(BatchRepeatLazyTensor(arg, batch_repeat=self.batch_repeat))
-            else:
-                args.append(arg)
-
-        new_lazy_tensor = self.base_lazy_tensor.__class__(*args, **kwargs)
-        return new_lazy_tensor.__getitem__(index)
diff --git a/gpytorch/lazy/block_lazy_tensor.py b/gpytorch/lazy/block_lazy_tensor.py
index 4daf8202..36fb4bc9 100644
--- a/gpytorch/lazy/block_lazy_tensor.py
+++ b/gpytorch/lazy/block_lazy_tensor.py
@@ -2,7 +2,6 @@
 
 import torch
 from .lazy_tensor import LazyTensor
-from .. import settings
 
 
 class BlockLazyTensor(LazyTensor):
@@ -38,57 +37,19 @@ class BlockLazyTensor(LazyTensor):
         self.base_lazy_tensor = base_lazy_tensor
         self.num_blocks = num_blocks
 
-    def _transpose_nonbatch(self):
-        return self.__class__(self.base_lazy_tensor._transpose_nonbatch(), num_blocks=self.num_blocks)
-
-    def mul(self, other):
-        # We're using a custom method here - the constant mul is applied to the base_lazy tensor
-        # This preserves the block structure
-
-        if not (torch.is_tensor(other) or isinstance(other, LazyTensor)) or (
-            torch.is_tensor(other) and other.numel() == 1
-        ):
-            from .constant_mul_lazy_tensor import ConstantMulLazyTensor
-
-            return self.__class__(ConstantMulLazyTensor(self.base_lazy_tensor, other), num_blocks=self.num_blocks)
-        else:
-            return super(BlockLazyTensor, self).mul(other)
-
-    def __getitem__(self, index):
-        # Process the index
-        ndimension = self.ndimension()
-        index = list(index) if isinstance(index, tuple) else [index]
-        index = [torch.tensor(idx) if isinstance(idx, list) else idx for idx in index]
-
-        # Handle the ellipsis
-        # Find the index of the ellipsis
-        ellipsis_locs = [index for index, item in enumerate(index) if item is Ellipsis]
-        if settings.debug.on():
-            if len(ellipsis_locs) > 1:
-                raise RuntimeError(
-                    "Cannot have multiple ellipsis in a __getitem__ call. LazyTensor {} "
-                    " received index {}.".format(self, index)
-                )
-        if len(ellipsis_locs) == 1:
-            ellipsis_loc = ellipsis_locs[0]
-            num_to_fill_in = ndimension - (len(index) - 1)
-            index = index[:ellipsis_loc] + [slice(None, None, None)] * num_to_fill_in + index[ellipsis_loc + 1:]
-
-        # Convert index back to tuple
-        index = tuple(index)
-
+    def _getitem(self, *indices):
         if self.num_blocks is None:
-            res = super(BlockLazyTensor, self).__getitem__(index)
+            res = super(BlockLazyTensor, self)._getitem(*indices)
             return res
 
         # Cases for when there's an inner batch
         else:
-            batch_index = index[0]
+            batch_index = indices[0]
             first_tensor_index_dim = None
 
             # Keeping all batch dimensions - recursion base case
             if isinstance(batch_index, slice) and batch_index == slice(None, None, None):
-                res = super(BlockLazyTensor, self).__getitem__(index)
+                res = super(BlockLazyTensor, self)._getitem(*indices)
                 return res
 
             # Construct a new lazy tensor
@@ -119,13 +80,13 @@ class BlockLazyTensor(LazyTensor):
             new_var = self.__class__(*components, num_blocks=num_blocks)
 
             # If the index was only on the batch index, we're done
-            if len(index) == 1:
+            if len(indices) == 1:
                 return new_var
 
             # Else - recurse
             else:
-                left_index = index[1]
-                right_index = index[2] if len(index) >= 3 else slice(None, None, None)
+                left_index = indices[1]
+                right_index = indices[2] if len(indices) >= 3 else slice(None, None, None)
 
                 # Normal case if we're indexing the LT with ints or slices
                 # Also squeeze dimensions if we're indexing with tensors
@@ -182,3 +143,19 @@ class BlockLazyTensor(LazyTensor):
                     res = res.squeeze(-1)
 
                 return res
+
+    def _transpose_nonbatch(self):
+        return self.__class__(self.base_lazy_tensor._transpose_nonbatch(), num_blocks=self.num_blocks)
+
+    def mul(self, other):
+        # We're using a custom method here - the constant mul is applied to the base_lazy tensor
+        # This preserves the block structure
+
+        if not (torch.is_tensor(other) or isinstance(other, LazyTensor)) or (
+            torch.is_tensor(other) and other.numel() == 1
+        ):
+            from .constant_mul_lazy_tensor import ConstantMulLazyTensor
+
+            return self.__class__(ConstantMulLazyTensor(self.base_lazy_tensor, other), num_blocks=self.num_blocks)
+        else:
+            return super(BlockLazyTensor, self).mul(other)
diff --git a/gpytorch/lazy/constant_mul_lazy_tensor.py b/gpytorch/lazy/constant_mul_lazy_tensor.py
index cafca1c0..0248c02d 100644
--- a/gpytorch/lazy/constant_mul_lazy_tensor.py
+++ b/gpytorch/lazy/constant_mul_lazy_tensor.py
@@ -72,15 +72,33 @@ class ConstantMulLazyTensor(LazyTensor):
         self.base_lazy_tensor = base_lazy_tensor
         self.constant = constant
 
-    def _matmul(self, rhs):
-        res = self.base_lazy_tensor._matmul(rhs)
+    def _approx_diag(self):
+        res = self.base_lazy_tensor._approx_diag()
         res_mat_shape = res.shape[len(self.base_lazy_tensor.batch_shape):]
         constant = self.constant.view(*self.constant.shape, *[1 for i in res_mat_shape])
-        res = res * constant
-        return res
+        return res * constant
 
-    def _t_matmul(self, rhs):
-        res = self.base_lazy_tensor._t_matmul(rhs)
+    def _get_indices(self, left_indices, right_indices, *batch_indices):
+        res = self.base_lazy_tensor._get_indices(left_indices, right_indices, *batch_indices)
+        constant = self.constant.__getitem__(batch_indices)
+        return res * constant
+
+    def _getitem(self, *indices):
+        # NOTE TO FUTURE SELF:
+        # This custom __getitem__ is actually very important!
+        # It prevents constructing an InterpolatedLazyTensor when one isn't needed
+        # This effects runntimes by up to 5x on simple exat GPs
+        # Run __getitem__ on the base_lazy_tensor and the constant
+        base_lazy_tensor = self.base_lazy_tensor._getitem(*indices)
+        constant = self.constant[indices[:-2]]
+
+        if torch.is_tensor(base_lazy_tensor):
+            constant = constant.view(*constant.shape, *[1] * (base_lazy_tensor.dim() - constant.dim()))
+
+        return base_lazy_tensor * constant
+
+    def _matmul(self, rhs):
+        res = self.base_lazy_tensor._matmul(rhs)
         res_mat_shape = res.shape[len(self.base_lazy_tensor.batch_shape):]
         constant = self.constant.view(*self.constant.shape, *[1 for i in res_mat_shape])
         res = res * constant
@@ -102,16 +120,18 @@ class ConstantMulLazyTensor(LazyTensor):
     def _size(self):
         return self.base_lazy_tensor.size()
 
+    def _t_matmul(self, rhs):
+        res = self.base_lazy_tensor._t_matmul(rhs)
+        res_mat_shape = res.shape[len(self.base_lazy_tensor.batch_shape):]
+        constant = self.constant.view(*self.constant.shape, *[1 for i in res_mat_shape])
+        res = res * constant
+        return res
+
     def _transpose_nonbatch(self):
         return ConstantMulLazyTensor(self.base_lazy_tensor._transpose_nonbatch(), self.constant)
 
-    def _get_indices(self, left_indices, right_indices, *batch_indices):
-        res = self.base_lazy_tensor._get_indices(left_indices, right_indices, *batch_indices)
-        constant = self.constant.__getitem__(batch_indices)
-        return res * constant
-
-    def _approx_diag(self):
-        res = self.base_lazy_tensor._approx_diag()
+    def diag(self):
+        res = self.base_lazy_tensor.diag()
         res_mat_shape = res.shape[len(self.base_lazy_tensor.batch_shape):]
         constant = self.constant.view(*self.constant.shape, *[1 for i in res_mat_shape])
         return res * constant
@@ -121,27 +141,3 @@ class ConstantMulLazyTensor(LazyTensor):
         res_mat_shape = res.shape[len(self.base_lazy_tensor.batch_shape):]
         constant = self.constant.view(*self.constant.shape, *[1 for i in res_mat_shape])
         return constant * res
-
-    def diag(self):
-        res = self.base_lazy_tensor.diag()
-        res_mat_shape = res.shape[len(self.base_lazy_tensor.batch_shape):]
-        constant = self.constant.view(*self.constant.shape, *[1 for i in res_mat_shape])
-        return res * constant
-
-    def __getitem__(self, i):
-        # TODO: Fix to support ellipsis
-        constant = self.constant
-        if constant.numel() > 1:
-            first_index = i[0] if isinstance(i, tuple) else i
-            constant = constant[first_index]
-
-        base_lazy_tensor = self.base_lazy_tensor.__getitem__(i)
-
-        if torch.is_tensor(base_lazy_tensor):
-            constant = constant.view(*constant.shape, *[1] * (len(base_lazy_tensor.shape) - len(constant.shape)))
-            return constant * base_lazy_tensor
-
-        if constant.numel() == 1:
-            constant = constant.squeeze()
-
-        return base_lazy_tensor * constant
diff --git a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
index 5fa313cd..cfdd779d 100644
--- a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
+++ b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
@@ -32,28 +32,84 @@ class LazyEvaluatedKernelTensor(LazyTensor):
     def device(self):
         return self.x1.device
 
-    def _matmul(self, rhs):
-        raise RuntimeError(LAZY_KERNEL_TENSOR_WARNING)
+    def _get_indices(self, left_indices, right_indices, *batch_indices):
+        from ..kernels import Kernel
 
-    def _t_matmul(self, rhs):
+        x1 = self.x1[(*batch_indices, left_indices)].unsqueeze(0)
+        x2 = self.x2[(*batch_indices, right_indices)].unsqueeze(0)
+        res = super(Kernel, self.kernel).__call__(x1.transpose(0, 1), x2.transpose(0, 1))
+        if isinstance(res, LazyTensor):
+            res = res.evaluate()
+        res = res.view(-1)
+        return res
+
+    def _getitem(self, *indices):
+        if self.is_batch:
+            batch_index = indices[0]
+            left_index = indices[1]
+            right_index = indices[2]
+            squeeze_row = self.squeeze_row
+            squeeze_col = self.squeeze_col
+
+            x1 = self.x1
+            if self.batch_dims == (0, 2):
+                x1 = x1.permute(0, 2, 1).contiguous()
+                x1 = x1.view(-1, x1.size(-1), 1)
+            x1 = x1[batch_index, left_index, :]
+            if x1.dim() == 2 and not isinstance(batch_index, int):
+                x1 = x1.unsqueeze(1)
+                squeeze_row = True
+            x2 = self.x2
+            if self.batch_dims == (0, 2):
+                x2 = x2.permute(0, 2, 1).contiguous()
+                x2 = x2.view(-1, x2.size(-1), 1)
+            x2 = x2[batch_index, right_index, :]
+            if x2.dim() == 2 and not isinstance(batch_index, int):
+                x2 = x2.unsqueeze(1)
+                squeeze_col = True
+
+            return LazyEvaluatedKernelTensor(
+                self.kernel, x1, x2, squeeze_row=squeeze_row, squeeze_col=squeeze_col, **self.params
+            )
+        else:
+            left_index = indices[0]
+            right_index = indices[1]
+            squeeze_row = self.squeeze_row
+            squeeze_col = self.squeeze_col
+
+            x1 = self.x1[left_index, :]
+            if x1.dim() == 1:
+                x1 = x1.unsqueeze(1)
+                squeeze_row = True
+            x2 = self.x2[right_index, :]
+            if x2.dim() == 1:
+                x2 = x2.unsqueeze(1)
+                squeeze_col = True
+
+            return LazyEvaluatedKernelTensor(
+                self.kernel, x1, x2, squeeze_row=squeeze_row, squeeze_col=squeeze_col, **self.params
+            )
+
+    def _matmul(self, rhs):
         raise RuntimeError(LAZY_KERNEL_TENSOR_WARNING)
 
     def _quad_form_derivative(self, left_vecs, right_vecs):
         raise RuntimeError(LAZY_KERNEL_TENSOR_WARNING)
 
+    def _size(self):
+        size = self.kernel.size(self.x1, self.x2)
+        if self.batch_dims == (0, 2):
+            if len(size) == 2:
+                return torch.Size((self.x1.size(-1), size[0], size[1]))
+            else:
+                return torch.Size((self.x1.size(-1) * size[0], size[1], size[2]))
+        return size
+
     def _transpose_nonbatch(self):
         return self.__class__(self.kernel, self.x2, self.x1, **self.params)
 
-    def _get_indices(self, left_indices, right_indices, *batch_indices):
-        from ..kernels import Kernel
-
-        x1 = self.x1.__getitem__((*batch_indices, left_indices)).unsqueeze(0)
-        x2 = self.x2.__getitem__((*batch_indices, right_indices)).unsqueeze(0)
-        res = super(Kernel, self.kernel).__call__(x1.transpose(0, 1), x2.transpose(0, 1))
-        if isinstance(res, LazyTensor):
-            res = res.evaluate()
-        res = res.view(-1)
-        return res
+    def _t_matmul(self, rhs):
+        raise RuntimeError(LAZY_KERNEL_TENSOR_WARNING)
 
     def diag(self):
         """
@@ -156,12 +212,6 @@ class LazyEvaluatedKernelTensor(LazyTensor):
                 self._cached_kernel_eval = NonLazyTensor(self._cached_kernel_eval)
             return self._cached_kernel_eval
 
-    def representation(self):
-        return self.evaluate_kernel().representation()
-
-    def representation_tree(self):
-        return LazyTensorRepresentationTree(self.evaluate_kernel())
-
     def evaluate(self):
         return self.evaluate_kernel().evaluate()
 
@@ -210,78 +260,8 @@ class LazyEvaluatedKernelTensor(LazyTensor):
 
         return LazyEvaluatedKernelTensor(self.kernel, x1, x2, **self.params)
 
-    def __getitem__(self, index):
-        ndimension = self.ndimension()
-        index = list(index) if isinstance(index, tuple) else [index]
-
-        # Handle the ellipsis
-        # Find the index of the ellipsis
-        ellipsis_locs = [index for index, item in enumerate(index) if item is Ellipsis]
-        if settings.debug.on():
-            if len(ellipsis_locs) > 1:
-                raise RuntimeError(
-                    "Cannot have multiple ellipsis in a __getitem__ call. LazyTensor {} "
-                    " received index {}.".format(self, index)
-                )
-        if len(ellipsis_locs) == 1:
-            ellipsis_loc = ellipsis_locs[0]
-            num_to_fill_in = ndimension - (len(index) - 1)
-            index = index[:ellipsis_loc] + [slice(None, None, None)] * num_to_fill_in + index[ellipsis_loc + 1:]
-
-        # Pad the index with empty slices
-        index += [slice(None, None, None)] * (ndimension - len(index))
-
-        if self.is_batch:
-            batch_index = index[0]
-            left_index = index[1]
-            right_index = index[2]
-            squeeze_row = self.squeeze_row
-            squeeze_col = self.squeeze_col
-
-            x1 = self.x1
-            if self.batch_dims == (0, 2):
-                x1 = x1.permute(0, 2, 1).contiguous()
-                x1 = x1.view(-1, x1.size(-1), 1)
-            x1 = x1[batch_index, left_index, :]
-            if x1.dim() == 2 and not isinstance(batch_index, int):
-                x1 = x1.unsqueeze(1)
-                squeeze_row = True
-            x2 = self.x2
-            if self.batch_dims == (0, 2):
-                x2 = x2.permute(0, 2, 1).contiguous()
-                x2 = x2.view(-1, x2.size(-1), 1)
-            x2 = x2[batch_index, right_index, :]
-            if x2.dim() == 2 and not isinstance(batch_index, int):
-                x2 = x2.unsqueeze(1)
-                squeeze_col = True
-
-            return LazyEvaluatedKernelTensor(
-                self.kernel, x1, x2, squeeze_row=squeeze_row, squeeze_col=squeeze_col, **self.params
-            )
-        else:
-            left_index = index[0]
-            right_index = index[1]
-            squeeze_row = self.squeeze_row
-            squeeze_col = self.squeeze_col
-
-            x1 = self.x1[left_index, :]
-            if x1.dim() == 1:
-                x1 = x1.unsqueeze(1)
-                squeeze_row = True
-            x2 = self.x2[right_index, :]
-            if x2.dim() == 1:
-                x2 = x2.unsqueeze(1)
-                squeeze_col = True
-
-            return LazyEvaluatedKernelTensor(
-                self.kernel, x1, x2, squeeze_row=squeeze_row, squeeze_col=squeeze_col, **self.params
-            )
+    def representation(self):
+        return self.evaluate_kernel().representation()
 
-    def _size(self):
-        size = self.kernel.size(self.x1, self.x2)
-        if self.batch_dims == (0, 2):
-            if len(size) == 2:
-                return torch.Size((self.x1.size(-1), size[0], size[1]))
-            else:
-                return torch.Size((self.x1.size(-1) * size[0], size[1], size[2]))
-        return size
+    def representation_tree(self):
+        return LazyTensorRepresentationTree(self.evaluate_kernel())
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index a2d34a58..5afa0a3e 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -63,10 +63,6 @@ class LazyTensor(object):
         operate on these batches if present.
     """
 
-    def __init__(self, *args, **kwargs):
-        self._args = args
-        self._kwargs = kwargs
-
     def _get_indices(self, *batch_indices, left_indices, right_indices):
         """
         Returns entries of the matrix, indexed by batch, row, and column indices
@@ -92,38 +88,6 @@ class LazyTensor(object):
         """
         raise NotImplementedError("The class {} requires a _matmul function!".format(self.__class__.__name__))
 
-    def _quad_form_derivative(self, left_vecs, right_vecs):
-        """
-        Given u (left_vecs) and v (right_vecs),
-        Computes the derivatives of (u^t K v) w.r.t. K
-
-        ..note::
-            This method is intended to be used only internally by various Functions that support backpropagation.
-            For example, this method is used internally by :func:`~gpytorch.lazy.LazyTensor.inv_quad_log_det`. It is
-            not likely that users will need to call this method directly.
-
-        Returns:
-            :obj:`torch.tensor`: derivative with respect to the arguments that are actually used to represent this
-                                   this LazyTensor.
-        """
-        args = self.representation()
-
-        toggled = [False] * len(args)
-
-        for i, arg in enumerate(args):
-            if not arg.requires_grad:
-                arg.requires_grad = True
-                toggled[i] = True
-
-        loss = (left_vecs * self._matmul(right_vecs)).sum()
-        grads = torch.autograd.grad(loss, args)
-
-        for i, arg in enumerate(args):
-            if toggled[i]:
-                arg.requires_grad = False
-
-        return grads
-
     def _size(self):
         """
         Returns the size of the resulting Tensor that the lazy tensor represents.
@@ -150,6 +114,10 @@ class LazyTensor(object):
             "The class {} requires a _transpose_nonbatch function!".format(self.__class__.__name__)
         )
 
+    def __init__(self, *args, **kwargs):
+        self._args = args
+        self._kwargs = kwargs
+
     @property
     def _args(self):
         return self._args_memo
@@ -203,6 +171,118 @@ class LazyTensor(object):
         # where S S^T = (K_XX + sigma^2 I)^-1
         return test_train_covar.matmul(precomputed_cache)
 
+    def _getitem(self, *indices):
+        """
+        Supports subindexing of the matrix this LazyTensor represents. This may return either another
+        :obj:`gpytorch.lazy.LazyTensor` or a :obj:`torch.tensor` depending on the exact implementation.
+
+        ..note::
+            LazyTensor.__getitem__ uses this as a helper method. If you are writing your own custom LazyTensor,
+            override this method rather than __getitem__ (so that you don't have to repeat the extra work)
+
+        ..note::
+            This method is used internally by the related function :func:`~gpytorch.lazy.LazyTensor.__getitem__`,
+            which does some additional work. Calling this method directly is discouraged.
+
+        Args:
+            :attr:`indices` (tuple of int, slice, or LongTensor):
+                A collection of indices for each of the dimensions. There will be exactly one index per dimension.
+        """
+        if settings.debug.on():
+            if len(indices) != self.dim():
+                raise RuntimeError(
+                    f"{self.__class__.__name__}._getitem() called with {len(indices)} indices - expected {self.dim()}. "
+                    "This is potentially a bug in GPyTorch."
+                )
+
+        components = list(self._args)
+        indices = list(indices)
+
+        # Normal case if we're indexing the LT with ints or slices
+        # Also squeeze dimensions if we're indexing with tensors
+        squeeze_left = False
+        squeeze_right = False
+        if isinstance(indices[-2], int):
+            indices[-2] = slice(indices[-2], indices[-2] + 1, None)
+            squeeze_left = True
+        elif torch.is_tensor(indices[-2]):
+            squeeze_left = True
+        if isinstance(indices[-1], int):
+            indices[-1] = slice(indices[-1], indices[-1] + 1, None)
+            squeeze_right = True
+        elif torch.is_tensor(indices[-1]):
+            squeeze_right = True
+
+        # Handle batch dimensions
+        isbatch = self.dim() >= 3
+        first_tensor_index_dim = None
+        if isbatch:
+            batch_index = tuple(indices[:-2])
+            for i, item in enumerate(components):
+                components[i] = item[batch_index]
+
+            for i, idx in enumerate(batch_index):
+                if torch.is_tensor(idx):
+                    first_tensor_index_dim = i
+                    break
+
+        new_lazy_tensor = self.__class__(*components, **self._kwargs)
+
+        # Handle index
+        left_index = indices[-2]
+        right_index = indices[-1]
+
+        # Special case: if both row and col are not indexed, then we are done
+        if (
+            not torch.is_tensor(left_index)
+            and left_index == slice(None, None, None)
+            and not torch.is_tensor(right_index)
+            and right_index == slice(None, None, None)
+        ):
+            return new_lazy_tensor
+
+        # Special case: if both row and col are tensor indexed, then we use _get_indices
+        if torch.is_tensor(left_index) and torch.is_tensor(right_index):
+            if left_index.numel() != right_index.numel():
+                raise RuntimeError(
+                    "Expected the tensor indices to be the same size: got {} and {}".format(
+                        left_index.numel(), right_index.numel()
+                    )
+                )
+
+            if new_lazy_tensor.ndimension() == 2:
+                return new_lazy_tensor._get_indices(left_index, right_index)
+
+            else:
+                batch_index = torch.arange(0, new_lazy_tensor.size(0), dtype=torch.long, device=self.device)
+                if first_tensor_index_dim is not None:
+                    if batch_index.numel() != left_index.numel():
+                        raise RuntimeError(
+                            "Expected the tensor indices to be the same size: got {}, {} and {}".format(
+                                batch_index.numel(), left_index.numel(), right_index.numel()
+                            )
+                        )
+                    return new_lazy_tensor._get_indices(left_index, right_index, batch_index)
+                else:
+                    batch_size = batch_index.numel()
+                    row_col_size = left_index.numel()
+                    batch_index = batch_index.unsqueeze(1).repeat(1, row_col_size).view(-1)
+                    left_index = left_index.unsqueeze(1).repeat(batch_size, 1).view(-1)
+                    right_index = right_index.unsqueeze(1).repeat(batch_size, 1).view(-1)
+                    res = new_lazy_tensor._get_indices(left_index, right_index, batch_index)
+                    return res.view(batch_size, row_col_size)
+
+        # Normal case: we have to do some processing on eithe rthe rows or columns
+        res = new_lazy_tensor._getitem_nonbatch(left_index, right_index, first_tensor_index_dim)
+        if (squeeze_left or squeeze_right) and isinstance(res, LazyTensor):
+            res = res.evaluate()
+        if squeeze_left:
+            res = res.squeeze(-2)
+        if squeeze_right:
+            res = res.squeeze(-1)
+
+        return res
+
     def _getitem_nonbatch(self, row_index, col_index, first_tensor_index_dim=None):
         """
         Given an index over rows and columns, gets those items from the LazyTensor.
@@ -292,6 +372,38 @@ class LazyTensor(object):
         else:
             return None
 
+    def _quad_form_derivative(self, left_vecs, right_vecs):
+        """
+        Given u (left_vecs) and v (right_vecs),
+        Computes the derivatives of (u^t K v) w.r.t. K
+
+        ..note::
+            This method is intended to be used only internally by various Functions that support backpropagation.
+            For example, this method is used internally by :func:`~gpytorch.lazy.LazyTensor.inv_quad_log_det`. It is
+            not likely that users will need to call this method directly.
+
+        Returns:
+            :obj:`torch.tensor`: derivative with respect to the arguments that are actually used to represent this
+                                   this LazyTensor.
+        """
+        args = self.representation()
+
+        toggled = [False] * len(args)
+
+        for i, arg in enumerate(args):
+            if not arg.requires_grad:
+                arg.requires_grad = True
+                toggled[i] = True
+
+        loss = (left_vecs * self._matmul(right_vecs)).sum()
+        grads = torch.autograd.grad(loss, args)
+
+        for i, arg in enumerate(args):
+            if toggled[i]:
+                arg.requires_grad = False
+
+        return grads
+
     def _preconditioner(self):
         """
         (Optional) define a preconditioner (P) for linear conjugate gradients
@@ -1239,15 +1351,14 @@ class LazyTensor(object):
         :obj:`gpytorch.lazy.LazyTensor` or a :obj:`torch.tensor` depending on the exact implementation.
         """
         ndimension = self.ndimension()
-        components = list(self._args)
 
         # Process the index
-        index = list(index) if isinstance(index, tuple) else [index]
-        index = [torch.tensor(idx) if isinstance(idx, list) else idx for idx in index]
+        index = index if isinstance(index, tuple) else (index,)
+        index = tuple(torch.tensor(idx) if isinstance(idx, list) else idx for idx in index)
 
         # Handle the ellipsis
         # Find the index of the ellipsis
-        ellipsis_locs = [index for index, item in enumerate(index) if item is Ellipsis]
+        ellipsis_locs = tuple(index for index, item in enumerate(index) if item is Ellipsis)
         if settings.debug.on():
             if len(ellipsis_locs) > 1:
                 raise RuntimeError(
@@ -1257,95 +1368,17 @@ class LazyTensor(object):
         if len(ellipsis_locs) == 1:
             ellipsis_loc = ellipsis_locs[0]
             num_to_fill_in = ndimension - (len(index) - 1)
-            index = index[:ellipsis_loc] + [slice(None, None, None)] * num_to_fill_in + index[ellipsis_loc + 1:]
+            index = index[:ellipsis_loc] + tuple(slice(None, None, None) for _ in range(num_to_fill_in)) + \
+                index[ellipsis_loc + 1:]
 
         # Pad the index with empty slices
-        index += [slice(None, None, None)] * (ndimension - len(index))
+        index = index + tuple(slice(None, None, None) for _ in range(ndimension - len(index)))
 
-        # Normal case if we're indexing the LT with ints or slices
-        # Also squeeze dimensions if we're indexing with tensors
-        squeeze_left = False
-        squeeze_right = False
-        if isinstance(index[-2], int):
-            index[-2] = slice(index[-2], index[-2] + 1, None)
-            squeeze_left = True
-        elif torch.is_tensor(index[-2]):
-            squeeze_left = True
-        if isinstance(index[-1], int):
-            index[-1] = slice(index[-1], index[-1] + 1, None)
-            squeeze_right = True
-        elif torch.is_tensor(index[-1]):
-            squeeze_right = True
-
-        # Handle batch dimensions
-        isbatch = ndimension >= 3
-        first_tensor_index_dim = None
-        if isbatch:
-            batch_index = tuple(index[:-2])
-            for i, item in enumerate(components):
-                components[i] = item[batch_index]
+        # Make the index a tuple again
+        index = tuple(index)
 
-            for i, idx in enumerate(batch_index):
-                if torch.is_tensor(idx):
-                    first_tensor_index_dim = i
-                    break
-
-        new_lazy_tensor = self.__class__(*components, **self._kwargs)
-
-        # Handle index
-        left_index = index[-2]
-        right_index = index[-1]
-
-        # Special case: if both row and col are not indexed, then we are done
-        if (
-            not torch.is_tensor(left_index)
-            and left_index == slice(None, None, None)
-            and not torch.is_tensor(right_index)
-            and right_index == slice(None, None, None)
-        ):
-            return new_lazy_tensor
-
-        # Special case: if both row and col are tensor indexed, then we use _get_indices
-        if torch.is_tensor(left_index) and torch.is_tensor(right_index):
-            if left_index.numel() != right_index.numel():
-                raise RuntimeError(
-                    "Expected the tensor indices to be the same size: got {} and {}".format(
-                        left_index.numel(), right_index.numel()
-                    )
-                )
-
-            if new_lazy_tensor.ndimension() == 2:
-                return new_lazy_tensor._get_indices(left_index, right_index)
-
-            else:
-                batch_index = torch.arange(0, new_lazy_tensor.size(0), dtype=torch.long, device=self.device)
-                if first_tensor_index_dim is not None:
-                    if batch_index.numel() != left_index.numel():
-                        raise RuntimeError(
-                            "Expected the tensor indices to be the same size: got {}, {} and {}".format(
-                                batch_index.numel(), left_index.numel(), right_index.numel()
-                            )
-                        )
-                    return new_lazy_tensor._get_indices(left_index, right_index, batch_index)
-                else:
-                    batch_size = batch_index.numel()
-                    row_col_size = left_index.numel()
-                    batch_index = batch_index.unsqueeze(1).repeat(1, row_col_size).view(-1)
-                    left_index = left_index.unsqueeze(1).repeat(batch_size, 1).view(-1)
-                    right_index = right_index.unsqueeze(1).repeat(batch_size, 1).view(-1)
-                    res = new_lazy_tensor._get_indices(left_index, right_index, batch_index)
-                    return res.view(batch_size, row_col_size)
-
-        # Normal case: we have to do some processing on eithe rthe rows or columns
-        res = new_lazy_tensor._getitem_nonbatch(left_index, right_index, first_tensor_index_dim)
-        if (squeeze_left or squeeze_right) and isinstance(res, LazyTensor):
-            res = res.evaluate()
-        if squeeze_left:
-            res = res.squeeze(-2)
-        if squeeze_right:
-            res = res.squeeze(-1)
-
-        return res
+        # Call self._getitem - now that the index has been processed
+        return self._getitem(*index)
 
     def __mul__(self, other):
         """
diff --git a/gpytorch/lazy/non_lazy_tensor.py b/gpytorch/lazy/non_lazy_tensor.py
index 845402b2..0f4a142e 100644
--- a/gpytorch/lazy/non_lazy_tensor.py
+++ b/gpytorch/lazy/non_lazy_tensor.py
@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 
 import torch
-from .. import settings
 from ..lazy import LazyTensor
 
 
@@ -19,12 +18,23 @@ class NonLazyTensor(LazyTensor):
         super(NonLazyTensor, self).__init__(tsr)
         self.tensor = tsr
 
+    def _get_indices(self, left_indices, right_indices, *batch_indices):
+        return self.tensor.__getitem__((*batch_indices, left_indices, right_indices))
+
+    def _getitem(self, *indices):
+        # Perform the __getitem__
+        res = self.tensor[indices]
+
+        row_index = indices[-2]
+        col_index = indices[-1]
+        if isinstance(row_index, int) or torch.is_tensor(row_index) or isinstance(col_index, int) or \
+                torch.is_tensor(col_index):
+            return res
+        return NonLazyTensor(res)
+
     def _matmul(self, rhs):
         return torch.matmul(self.tensor, rhs)
 
-    def _t_matmul(self, rhs):
-        return torch.matmul(self.tensor.transpose(-1, -2), rhs)
-
     def _quad_form_derivative(self, left_vecs, right_vecs):
         res = left_vecs.matmul(right_vecs.transpose(-1, -2))
         return (res,)
@@ -35,8 +45,8 @@ class NonLazyTensor(LazyTensor):
     def _transpose_nonbatch(self):
         return NonLazyTensor(self.tensor.transpose(-1, -2))
 
-    def _get_indices(self, left_indices, right_indices, *batch_indices):
-        return self.tensor.__getitem__((*batch_indices, left_indices, right_indices))
+    def _t_matmul(self, rhs):
+        return torch.matmul(self.tensor.transpose(-1, -2), rhs)
 
     def diag(self):
         if self.tensor.ndimension() < 3:
@@ -47,41 +57,3 @@ class NonLazyTensor(LazyTensor):
 
     def evaluate(self):
         return self.tensor
-
-    def __getitem__(self, index):
-        # Make sure index is a list
-        if not isinstance(index, tuple):
-            index = [index]
-        else:
-            index = list(index)
-
-        # Handle the ellipsis
-        # Find the index of the ellipsis
-        ellipsis_locs = [index for index, item in enumerate(index) if item is Ellipsis]
-        if settings.debug.on():
-            if len(ellipsis_locs) > 1:
-                raise RuntimeError(
-                    "Cannot have multiple ellipsis in a __getitem__ call. LazyTensor {} "
-                    " received index {}.".format(self, index)
-                )
-        if len(ellipsis_locs) == 1:
-            ellipsis_loc = ellipsis_locs[0]
-            num_to_fill_in = self.ndimension() - (len(index) - 1)
-            index = index[:ellipsis_loc] + [slice(None, None, None)] * num_to_fill_in + index[ellipsis_loc + 1:]
-
-        # Make the index a tuple again
-        index = tuple(index)
-
-        # Perform the __getitem__
-        res = self.tensor[index]
-
-        # Determine if we should return a LazyTensor or a Tensor
-        if len(index) >= self.ndimension() - 1:
-            row_index = index[self.ndimension() - 2]
-            if isinstance(row_index, int) or torch.is_tensor(row_index):
-                return res
-        if len(index) == self.ndimension():
-            col_index = index[self.ndimension() - 1]
-            if isinstance(col_index, int) or torch.is_tensor(col_index):
-                return res
-        return NonLazyTensor(res)
diff --git a/gpytorch/lazy/sum_lazy_tensor.py b/gpytorch/lazy/sum_lazy_tensor.py
index d2cd4379..d47acc09 100644
--- a/gpytorch/lazy/sum_lazy_tensor.py
+++ b/gpytorch/lazy/sum_lazy_tensor.py
@@ -19,12 +19,22 @@ class SumLazyTensor(LazyTensor):
 
         self.lazy_tensors = lazy_tensors
 
+    def _get_indices(self, left_indices, right_indices, *batch_indices):
+        return sum(
+            lazy_tensor._get_indices(left_indices, right_indices, *batch_indices)
+            for lazy_tensor in self.lazy_tensors
+        )
+
+    def _getitem(self, *indices):
+        results = tuple(lazy_tensor._getitem(*indices) for lazy_tensor in self.lazy_tensors)
+        if isinstance(results[0], LazyTensor):
+            return SumLazyTensor(*results)
+        else:
+            return sum(results)
+
     def _matmul(self, rhs):
         return sum(lazy_tensor._matmul(rhs) for lazy_tensor in self.lazy_tensors)
 
-    def _t_matmul(self, rhs):
-        return sum(lazy_tensor._t_matmul(rhs) for lazy_tensor in self.lazy_tensors)
-
     def _quad_form_derivative(self, left_vecs, right_vecs):
         return tuple(
             var for lazy_tensor in self.lazy_tensors for var in lazy_tensor._quad_form_derivative(left_vecs, right_vecs)
@@ -33,16 +43,13 @@ class SumLazyTensor(LazyTensor):
     def _size(self):
         return self.lazy_tensors[0].size()
 
+    def _t_matmul(self, rhs):
+        return sum(lazy_tensor._t_matmul(rhs) for lazy_tensor in self.lazy_tensors)
+
     def _transpose_nonbatch(self):
         lazy_tensors_t = [lazy_tensor.transpose(-1, -2) for lazy_tensor in self.lazy_tensors]
         return SumLazyTensor(*lazy_tensors_t)
 
-    def _get_indices(self, left_indices, right_indices, *batch_indices):
-        return sum(
-            lazy_tensor._get_indices(left_indices, right_indices, *batch_indices)
-            for lazy_tensor in self.lazy_tensors
-        )
-
     def _exact_predictive_covar_inv_quad_form_cache(self, train_train_covar_inv_root, test_train_covar):
         return tuple(
             lazy_tensor._exact_predictive_covar_inv_quad_form_cache(
@@ -88,10 +95,3 @@ class SumLazyTensor(LazyTensor):
 
     def sum_batch(self, sum_batch_size=None):
         return self.__class__(*(lazy_tensor.sum_batch(sum_batch_size) for lazy_tensor in self.lazy_tensors))
-
-    def __getitem__(self, index):
-        results = tuple(lazy_tensor.__getitem__(index) for lazy_tensor in self.lazy_tensors)
-        if isinstance(results[0], LazyTensor):
-            return SumLazyTensor(*results)
-        else:
-            return sum(results)
diff --git a/gpytorch/lazy/zero_lazy_tensor.py b/gpytorch/lazy/zero_lazy_tensor.py
index 7e26c5fe..047e3519 100644
--- a/gpytorch/lazy/zero_lazy_tensor.py
+++ b/gpytorch/lazy/zero_lazy_tensor.py
@@ -2,7 +2,6 @@
 
 import torch
 from .lazy_tensor import LazyTensor
-from .. import settings
 
 
 class ZeroLazyTensor(LazyTensor):
@@ -25,13 +24,35 @@ class ZeroLazyTensor(LazyTensor):
     def device(self):
         return self._device
 
-    def _matmul(self, rhs):
-        rhs_size_ind = -2 if rhs.ndimension() > 1 else -1
-        if self.size(-1) != rhs.size(rhs_size_ind):
-            raise RuntimeError("Size mismatch, self: {}, rhs: {}".format(self.size(), rhs.size()))
-        return rhs * 0
+    def _getitem(self, *indices):
+        has_added_tensor_index = False
+        evaluate = False
+        new_sizes = []
 
-    def _t_matmul(self, rhs):
+        for ix, sub_index in enumerate(indices):
+            if isinstance(sub_index, int):
+                if ix >= self.dim() - 2:
+                    evaluate = True
+                continue
+
+            elif torch.is_tensor(sub_index):
+                if ix >= self.dim() - 2:
+                    evaluate = True
+                if not has_added_tensor_index:
+                    new_sizes.append(sub_index.numel())
+                    has_added_tensor_index = True
+
+            elif isinstance(sub_index, slice):
+                new_sizes.append(len(torch.arange(self.size(ix))[sub_index]))
+            else:
+                new_sizes.append(len(sub_index))
+
+        if evaluate:
+            return torch.zeros(*new_sizes, dtype=self.dtype, device=self.device)
+        else:
+            return ZeroLazyTensor(*new_sizes)
+
+    def _matmul(self, rhs):
         rhs_size_ind = -2 if rhs.ndimension() > 1 else -1
         if self.size(-1) != rhs.size(rhs_size_ind):
             raise RuntimeError("Size mismatch, self: {}, rhs: {}".format(self.size(), rhs.size()))
@@ -43,6 +64,12 @@ class ZeroLazyTensor(LazyTensor):
     def _size(self):
         return torch.Size(self.sizes)
 
+    def _t_matmul(self, rhs):
+        rhs_size_ind = -2 if rhs.ndimension() > 1 else -1
+        if self.size(-1) != rhs.size(rhs_size_ind):
+            raise RuntimeError("Size mismatch, self: {}, rhs: {}".format(self.size(), rhs.size()))
+        return rhs * 0
+
     def add_diag(self, diag):
         from .diag_lazy_tensor import DiagLazyTensor
 
@@ -143,50 +170,3 @@ class ZeroLazyTensor(LazyTensor):
 
     def __mul__(self, other):
         return self
-
-    def __getitem__(self, index):
-        ndimension = self.ndimension()
-        index = list(index) if isinstance(index, tuple) else [index]
-
-        # Handle the ellipsis
-        # Find the index of the ellipsis
-        ellipsis_locs = [index for index, item in enumerate(index) if item is Ellipsis]
-        if settings.debug.on():
-            if len(ellipsis_locs) > 1:
-                raise RuntimeError(
-                    "Cannot have multiple ellipsis in a __getitem__ call. LazyTensor {} "
-                    " received index {}.".format(self, index)
-                )
-        if len(ellipsis_locs) == 1:
-            ellipsis_loc = ellipsis_locs[0]
-            num_to_fill_in = ndimension - (len(index) - 1)
-            index = index[:ellipsis_loc] + [slice(None, None, None)] * num_to_fill_in + index[ellipsis_loc + 1:]
-
-        # Pad the index with empty slices
-        index += [slice(None, None, None)] * (ndimension - len(index))
-
-        has_added_tensor_index = False
-        evaluate = False
-        new_sizes = []
-        for ix, sub_index in enumerate(index):
-            if isinstance(sub_index, int):
-                if ix >= ndimension - 2:
-                    evaluate = True
-                continue
-
-            elif torch.is_tensor(sub_index):
-                if ix >= ndimension - 2:
-                    evaluate = True
-                if not has_added_tensor_index:
-                    new_sizes.append(sub_index.numel())
-                    has_added_tensor_index = True
-
-            elif isinstance(sub_index, slice):
-                new_sizes.append(len(torch.arange(self.size(ix))[sub_index]))
-            else:
-                new_sizes.append(len(sub_index))
-
-        if evaluate:
-            return torch.zeros(*new_sizes, dtype=self.dtype, device=self.device)
-        else:
-            return ZeroLazyTensor(*new_sizes)


Test patch:
diff --git a/test/lazy/test_zero_lazy_tensor.py b/test/lazy/test_zero_lazy_tensor.py
index d2dbec4d..7bcf272f 100644
--- a/test/lazy/test_zero_lazy_tensor.py
+++ b/test/lazy/test_zero_lazy_tensor.py
@@ -28,8 +28,8 @@ class TestZeroLazyTensor(unittest.TestCase):
         lv = ZeroLazyTensor(5, 4, 3)
 
         res_one = lv[[0, 1]].evaluate()
-        res_two = lv[:, [0, 1], :].evaluate()
-        res_three = lv[:, :, [0, 2]].evaluate()
+        res_two = lv[:, [0, 1], :]
+        res_three = lv[:, :, [0, 2]]
 
         self.assertLess(torch.norm(res_one - torch.zeros(2, 4, 3)), 1e-4)
         self.assertLess(torch.norm(res_two - torch.zeros(5, 2, 3)), 1e-4)
@@ -39,8 +39,8 @@ class TestZeroLazyTensor(unittest.TestCase):
         lv = ZeroLazyTensor(5, 4, 3)
 
         res_one = lv[[0, 1]].evaluate()
-        res_two = lv[:, [0, 1], ...].evaluate()
-        res_three = lv[..., [0, 2]].evaluate()
+        res_two = lv[:, [0, 1], ...]
+        res_three = lv[..., [0, 2]]
 
         self.assertLess(torch.norm(res_one - torch.zeros(2, 4, 3)), 1e-4)
         self.assertLess(torch.norm(res_two - torch.zeros(5, 2, 3)), 1e-4)