Repository: cornellius-gp/gpytorch
Pull request #397 (2018-11-26 20:12:24)

LazyTensor.__getitem__ supports ellipsis
**MERGE #394 FIRST**

We can now write calls like `lazy_tensor[..., 4:, 3]`. This reduces the amount of casing we do checking for batch dimensions. It will also make it much easier to write multi-batch functions for LazyTensors.

A step towards #369 


Patch:
diff --git a/gpytorch/lazy/batch_repeat_lazy_tensor.py b/gpytorch/lazy/batch_repeat_lazy_tensor.py
index 4b85dab3..d2400fcd 100644
--- a/gpytorch/lazy/batch_repeat_lazy_tensor.py
+++ b/gpytorch/lazy/batch_repeat_lazy_tensor.py
@@ -18,17 +18,11 @@ class BatchRepeatLazyTensor(LazyTensor):
         self.base_lazy_tensor = base_lazy_tensor
         self.batch_repeat = batch_repeat
 
-    def _batch_get_indices(self, batch_indices, row_indices, col_indices):
-        return self._get_indices(batch_indices, row_indices=row_indices, col_indices=col_indices)
-
-    def _get_indices(self, *batch_indices, row_indices, col_indices):
+    def _get_indices(self, row_indices, col_indices, *batch_indices):
         num_true_batch_dims = len(self.base_lazy_tensor.batch_shape)
         batch_indices = [index % size for index, size in zip(batch_indices, self._padded_base_batch_shape)]
         batch_indices = batch_indices[-num_true_batch_dims:] if num_true_batch_dims else []
-        if len(batch_indices) == 1:
-            return self.base_lazy_tensor._batch_get_indices(*batch_indices, row_indices, col_indices)
-        else:
-            return self.base_lazy_tensor._get_indices(*batch_indices, row_indices, col_indices)
+        return self.base_lazy_tensor._get_indices(row_indices, col_indices, *batch_indices)
 
     def _matmul(self, rhs):
         rhs = self._move_repeat_batches_to_columns(rhs)
@@ -175,10 +169,6 @@ class BatchRepeatLazyTensor(LazyTensor):
         )
 
     def __getitem__(self, index):
-        """
-        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.
-        """
         args = []
         kwargs = self.base_lazy_tensor._kwargs
         num_base_batch_dims = len(self.base_lazy_tensor.batch_shape)
diff --git a/gpytorch/lazy/block_diag_lazy_tensor.py b/gpytorch/lazy/block_diag_lazy_tensor.py
index 2c14dd86..b2ac46b8 100644
--- a/gpytorch/lazy/block_diag_lazy_tensor.py
+++ b/gpytorch/lazy/block_diag_lazy_tensor.py
@@ -2,6 +2,7 @@
 
 import torch
 from .block_lazy_tensor import BlockLazyTensor
+from .. import settings
 
 
 class BlockDiagLazyTensor(BlockLazyTensor):
@@ -66,32 +67,37 @@ class BlockDiagLazyTensor(BlockLazyTensor):
             true_batch_size = self.base_lazy_tensor.size(0) // self.num_blocks
             return torch.Size((true_batch_size, self.num_blocks * base_size[1], self.num_blocks * base_size[2]))
 
-    def _batch_get_indices(self, batch_indices, left_indices, right_indices):
-        block_size = self.base_lazy_tensor.size(-1)
-        left_batch_indices = left_indices.div(block_size).long()
-        right_batch_indices = left_indices.div(block_size).long()
-        batch_indices = batch_indices * block_size + left_batch_indices
-        left_indices = left_indices.fmod(block_size)
-        right_indices = right_indices.fmod(block_size)
-
-        res = self.base_lazy_tensor._batch_get_indices(batch_indices, left_indices, right_indices)
-        res = res * torch.eq(left_batch_indices, right_batch_indices).type_as(res)
-        return res
-
-    def _get_indices(self, left_indices, right_indices):
-        block_size = self.base_lazy_tensor.size(-1)
-        left_batch_indices = left_indices.div(block_size).long()
-        right_batch_indices = left_indices.div(block_size).long()
-        left_indices = left_indices.fmod(block_size)
-        right_indices = right_indices.fmod(block_size)
-
-        res = self.base_lazy_tensor._batch_get_indices(left_batch_indices, left_indices, right_indices)
-        res = res * torch.eq(left_batch_indices, right_batch_indices).type_as(res)
-        return res
+    def _get_indices(self, left_indices, right_indices, *batch_indices):
+        if self.num_blocks is None:
+            if settings.debug.on():
+                assert(len(batch_indices) == 0)
+            block_size = self.base_lazy_tensor.size(-1)
+            left_batch_indices = left_indices.div(block_size).long()
+            right_batch_indices = left_indices.div(block_size).long()
+            left_indices = left_indices.fmod(block_size)
+            right_indices = right_indices.fmod(block_size)
+
+            res = self.base_lazy_tensor._get_indices(left_indices, right_indices, left_batch_indices)
+            res = res * torch.eq(left_batch_indices, right_batch_indices).type_as(res)
+            return res
+        else:
+            if settings.debug.on():
+                assert(len(batch_indices) == 1)
+            batch_indices = batch_indices[0]
+            block_size = self.base_lazy_tensor.size(-1)
+            left_batch_indices = left_indices.div(block_size).long()
+            right_batch_indices = left_indices.div(block_size).long()
+            batch_indices = batch_indices * block_size + left_batch_indices
+            left_indices = left_indices.fmod(block_size)
+            right_indices = right_indices.fmod(block_size)
+
+            res = self.base_lazy_tensor._get_indices(left_indices, right_indices, batch_indices)
+            res = res * torch.eq(left_batch_indices, right_batch_indices).type_as(res)
+            return res
 
     def diag(self):
         res = self.base_lazy_tensor.diag().contiguous()
-        if self.num_blocks:
+        if self.num_blocks is not None:
             res = res.view(res.size(0) // self.num_blocks, -1)
         else:
             res = res.view(-1)
diff --git a/gpytorch/lazy/block_lazy_tensor.py b/gpytorch/lazy/block_lazy_tensor.py
index e62d8bde..4daf8202 100644
--- a/gpytorch/lazy/block_lazy_tensor.py
+++ b/gpytorch/lazy/block_lazy_tensor.py
@@ -2,6 +2,7 @@
 
 import torch
 from .lazy_tensor import LazyTensor
+from .. import settings
 
 
 class BlockLazyTensor(LazyTensor):
@@ -54,13 +55,35 @@ class BlockLazyTensor(LazyTensor):
             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)
+
         if self.num_blocks is None:
             res = super(BlockLazyTensor, self).__getitem__(index)
             return res
 
         # Cases for when there's an inner batch
         else:
-            batch_index = index if not isinstance(index, tuple) else index[0]
+            batch_index = index[0]
             first_tensor_index_dim = None
 
             # Keeping all batch dimensions - recursion base case
@@ -96,7 +119,7 @@ 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 not isinstance(index, tuple) or len(index) == 1:
+            if len(index) == 1:
                 return new_var
 
             # Else - recurse
@@ -139,14 +162,14 @@ class BlockLazyTensor(LazyTensor):
                                         batch_index.numel(), left_index.numel(), right_index.numel()
                                     )
                                 )
-                            return new_var._batch_get_indices(batch_index, left_index, right_index)
+                            return new_var._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_var._batch_get_indices(batch_index, left_index, right_index)
+                            res = new_var._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
diff --git a/gpytorch/lazy/constant_mul_lazy_tensor.py b/gpytorch/lazy/constant_mul_lazy_tensor.py
index e54f3158..e221c2c2 100644
--- a/gpytorch/lazy/constant_mul_lazy_tensor.py
+++ b/gpytorch/lazy/constant_mul_lazy_tensor.py
@@ -54,24 +54,19 @@ class ConstantMulLazyTensor(LazyTensor):
     """
 
     def __init__(self, base_lazy_tensor, constant):
-        if torch.is_tensor(constant):
-            if constant.ndimension() > 1:
-                raise RuntimeError(
-                    "Got a constant with %d dimensions - expected a 0D or 1D tensor" % constant.ndimension()
-                )
-            elif constant.numel() > 1:
-                if not (base_lazy_tensor.ndimension() == 3 and base_lazy_tensor.size(0) == constant.numel()):
-                    numel = constant.numel()
-                    raise RuntimeError(
-                        "A constant with size %d expedts a 3D lazy var. with batch size %d. "
-                        "Got a %dD lazy var. with size %s"
-                        % (numel, numel, base_lazy_tensor.ndimension(), repr(base_lazy_tensor.size()))
-                    )
-
-            elif constant.numel() == 1:
-                constant = constant.squeeze()
-        else:
+        if not torch.is_tensor(constant):
             constant = torch.tensor(constant, device=base_lazy_tensor.device, dtype=base_lazy_tensor.dtype)
+        orig_constant = constant
+
+        # Make sure that the constant can be expanded to the appropriate size
+        try:
+            constant = orig_constant.expand(base_lazy_tensor.batch_shape)
+        except RuntimeError:
+            raise RuntimeError(
+                "ConstantMulLazyTensor of size {} received an invalid constant of size {}.".format(
+                    base_lazy_tensor.shape, orig_constant.shape
+                )
+            )
 
         super(ConstantMulLazyTensor, self).__init__(base_lazy_tensor, constant)
         self.base_lazy_tensor = base_lazy_tensor
@@ -79,39 +74,27 @@ class ConstantMulLazyTensor(LazyTensor):
 
     def _matmul(self, rhs):
         res = self.base_lazy_tensor._matmul(rhs)
-        res = res * self._constant_as(res)
+        constant = self.constant.view(*self.constant.shape, 1, 1)
+        res = res * constant
         return res
 
     def _t_matmul(self, rhs):
         res = self.base_lazy_tensor._t_matmul(rhs)
-        res = res * self._constant_as(res)
+        constant = self.constant.view(*self.constant.shape, 1, 1)
+        res = res * constant
         return res
 
     def _quad_form_derivative(self, left_vecs, right_vecs):
-        res = list(self.base_lazy_tensor._quad_form_derivative(left_vecs, right_vecs))
-        for i, item in enumerate(res):
-            if torch.is_tensor(item) and res[i].sum():
-                res[i] = res[i] * self._constant_as(res[i])
         # Gradient with respect to the constant
-        if self.constant.numel() == 1:
-            constant_deriv = (left_vecs * self.base_lazy_tensor._matmul(right_vecs)).sum().expand_as(self.constant)
-        else:
-            constant_deriv = left_vecs * self.base_lazy_tensor._matmul(right_vecs)
-            constant_deriv = constant_deriv.sum(-2, keepdim=True).sum(-1, keepdim=True)
-
-        if constant_deriv.dim():
-            constant_deriv = constant_deriv.view(*self.constant.size())
-        res.append(constant_deriv)
-        return res
-
-    def _constant_as(self, other):
-        size = [self.constant.numel()] + [1] * (other.ndimension() - 1)
-        constant = self.constant.view(*size)
+        constant_deriv = left_vecs * self.base_lazy_tensor._matmul(right_vecs)
+        constant_deriv = constant_deriv.sum(-2).sum(-1)
 
-        if constant.ndimension() > other.ndimension():
-            constant = constant.squeeze(-1)
+        # Get derivaties of everything else
+        constant = self.constant.view(*self.constant.shape, 1, 1)
+        left_vecs = left_vecs * constant
+        res = self.base_lazy_tensor._quad_form_derivative(left_vecs, right_vecs)
 
-        return constant
+        return res + (constant_deriv,)
 
     def _size(self):
         return self.base_lazy_tensor.size()
@@ -119,33 +102,22 @@ class ConstantMulLazyTensor(LazyTensor):
     def _transpose_nonbatch(self):
         return ConstantMulLazyTensor(self.base_lazy_tensor._transpose_nonbatch(), self.constant)
 
-    def _batch_get_indices(self, batch_indices, left_indices, right_indices):
-        res = self.base_lazy_tensor._batch_get_indices(batch_indices, left_indices, right_indices)
-        return self.constant.expand_as(res) * res
-
-    def _get_indices(self, left_indices, right_indices):
-        res = self.base_lazy_tensor._get_indices(left_indices, right_indices)
-        return self.constant.expand_as(res) * res
+    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()
-        return res * self._constant_as(res)
+        constant = self.constant.view(*self.constant.shape, 1)
+        return res * constant
 
     def evaluate(self):
         res = self.base_lazy_tensor.evaluate()
-        return res * self._constant_as(res)
+        constant = self.constant.view(*self.constant.shape, 1, 1)
+        return res * constant
 
     def diag(self):
         res = self.base_lazy_tensor.diag()
-        res = res * self._constant_as(res)
-        return res
-
-    def __getitem__(self, i):
-        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) and constant.dim() < base_lazy_tensor.dim():
-            constant = constant.view(constant.numel(), *([1] * (base_lazy_tensor.dim() - 1)))
-        return base_lazy_tensor * constant
+        constant = self.constant.view(*self.constant.shape, 1)
+        return res * constant
diff --git a/gpytorch/lazy/diag_lazy_tensor.py b/gpytorch/lazy/diag_lazy_tensor.py
index 58dd62ef..7326b662 100644
--- a/gpytorch/lazy/diag_lazy_tensor.py
+++ b/gpytorch/lazy/diag_lazy_tensor.py
@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 
 from itertools import product
-import warnings
 import torch
 from .lazy_tensor import LazyTensor
 from .root_lazy_tensor import RootLazyTensor
@@ -28,10 +27,6 @@ class DiagLazyTensor(LazyTensor):
 
         return AddedDiagLazyTensor(other, self)
 
-    def _batch_get_indices(self, batch_indices, left_indices, right_indices):
-        warnings.warn("_batch_get_indices is deprecated - use _get_indices instead", DeprecationWarning)
-        return self._get_indices(left_indices, right_indices, batch_indices)
-
     def _get_indices(self, left_indices, right_indices, *batch_indices):
         """Extract elements from the LazyTensor. Supports arbitrary batch sizes.
 
diff --git a/gpytorch/lazy/interpolated_lazy_tensor.py b/gpytorch/lazy/interpolated_lazy_tensor.py
index f36e045e..2109fa01 100644
--- a/gpytorch/lazy/interpolated_lazy_tensor.py
+++ b/gpytorch/lazy/interpolated_lazy_tensor.py
@@ -232,11 +232,11 @@ class InterpolatedLazyTensor(LazyTensor):
         )
         return res
 
-    def _batch_get_indices(self, batch_indices, left_indices, right_indices):
-        left_interp_indices = self.left_interp_indices[batch_indices, left_indices, :]
-        left_interp_values = self.left_interp_values[batch_indices, left_indices, :]
-        right_interp_indices = self.right_interp_indices[batch_indices, right_indices, :]
-        right_interp_values = self.right_interp_values[batch_indices, right_indices, :]
+    def _get_indices(self, left_indices, right_indices, *batch_indices):
+        left_interp_indices = self.left_interp_indices.__getitem__((*batch_indices, left_indices))
+        left_interp_values = self.left_interp_values.__getitem__((*batch_indices, left_indices))
+        right_interp_indices = self.right_interp_indices.__getitem__((*batch_indices, right_indices))
+        right_interp_values = self.right_interp_values.__getitem__((*batch_indices, right_indices))
 
         n_data, n_interp = left_interp_indices.size()
 
@@ -246,43 +246,19 @@ class InterpolatedLazyTensor(LazyTensor):
         interp_values = torch.matmul(left_interp_values, right_interp_values)
 
         # Batch compute values that will be non-zero for row k
-        left_interp_indices = left_interp_indices.unsqueeze(-1).expand(n_data, n_interp, n_interp)
-        right_interp_indices = right_interp_indices.unsqueeze(-2).expand(n_data, n_interp, n_interp)
-        left_interp_indices = left_interp_indices.contiguous()
-        right_interp_indices = right_interp_indices.contiguous()
-        batch_indices = batch_indices.unsqueeze(1).repeat(1, n_interp ** 2).view(-1)
-        base_var_vals = self.base_lazy_tensor._batch_get_indices(
-            batch_indices, left_interp_indices.view(-1), right_interp_indices.view(-1)
+        left_interp_indices = left_interp_indices.unsqueeze(-1).expand(n_data, n_interp, n_interp).contiguous()
+        right_interp_indices = right_interp_indices.unsqueeze(-2).expand(n_data, n_interp, n_interp).contiguous()
+        batch_indices = [
+            batch_index.unsqueeze(1).repeat(1, n_interp ** 2).view(-1)
+            for batch_index in batch_indices
+        ]
+        base_var_vals = self.base_lazy_tensor._get_indices(
+            left_interp_indices.view(-1), right_interp_indices.view(-1), *batch_indices
         )
         base_var_vals = base_var_vals.view(left_interp_indices.size())
         res = (interp_values * base_var_vals).sum(-1).sum(-1)
         return res
 
-    def _get_indices(self, left_indices, right_indices):
-        left_interp_indices = self.left_interp_indices[left_indices, :]
-        left_interp_values = self.left_interp_values[left_indices, :]
-        right_interp_indices = self.right_interp_indices[right_indices, :]
-        right_interp_values = self.right_interp_values[right_indices, :]
-
-        n_data, n_interp = left_interp_indices.size()
-
-        # Batch compute the non-zero values of the outer products w_left^k w_right^k^T
-        left_interp_values = left_interp_values.unsqueeze(-1)
-        right_interp_values = right_interp_values.unsqueeze(-2)
-        interp_values = torch.matmul(left_interp_values, right_interp_values)
-
-        # Batch compute values that will be non-zero for row k
-        if left_interp_indices.ndimension() == 3:
-            left_interp_indices = left_interp_indices.unsqueeze(-1).expand(n_data, n_interp, n_interp).contiguous()
-            right_interp_indices = right_interp_indices.unsqueeze(-2).expand(n_data, n_interp, n_interp).contiguous()
-        else:
-            left_interp_indices = left_interp_indices.unsqueeze(-1).expand(n_data, n_interp, n_interp).contiguous()
-            right_interp_indices = right_interp_indices.unsqueeze(-2).expand(n_data, n_interp, n_interp).contiguous()
-        base_var_vals = self.base_lazy_tensor._get_indices(left_interp_indices.view(-1), right_interp_indices.view(-1))
-        base_var_vals = base_var_vals.view(left_interp_indices.size())
-        res = (interp_values * base_var_vals).sum(-1).sum(-1)
-        return res
-
     def _sparse_left_interp_t(self, left_interp_indices_tensor, left_interp_values_tensor):
         if hasattr(self, "_sparse_left_interp_t_memo"):
             if torch.equal(self._left_interp_indices_memo, left_interp_indices_tensor) and torch.equal(
@@ -346,16 +322,16 @@ class InterpolatedLazyTensor(LazyTensor):
             left_interp_indices = left_interp_indices.contiguous().view(-1)
             right_interp_indices = right_interp_indices.contiguous().view(-1)
 
-            if self.base_lazy_tensor.ndimension() == 3:
-                batch_indices = torch.arange(0, batch_size, dtype=torch.long, device=self.device).unsqueeze(-1)
-                batch_indices = batch_indices.repeat(1, n_data * n_interp * n_interp).view(-1)
-                base_var_vals = self.base_lazy_tensor._batch_get_indices(
-                    batch_indices, left_interp_indices, right_interp_indices
-                )
-            else:
-                base_var_vals = self.base_lazy_tensor._get_indices(left_interp_indices, right_interp_indices)
-
+            base_batch_indices = [
+                torch.arange(0, batch_size, dtype=torch.long, device=self.device).
+                unsqueeze(-1).
+                repeat(1, n_data * n_interp * n_interp).view(-1)
+            ] if batch_size is not None else []
+            base_var_vals = self.base_lazy_tensor._get_indices(
+                left_interp_indices, right_interp_indices, *base_batch_indices
+            )
             res = interp_values * base_var_vals
+
             if batch_size is None:
                 res = res.view(n_data, -1).sum(-1)
             else:
diff --git a/gpytorch/lazy/kronecker_product_lazy_tensor.py b/gpytorch/lazy/kronecker_product_lazy_tensor.py
index 8fa2ba2e..c3ea8c9a 100644
--- a/gpytorch/lazy/kronecker_product_lazy_tensor.py
+++ b/gpytorch/lazy/kronecker_product_lazy_tensor.py
@@ -193,7 +193,7 @@ class KroneckerProductLazyTensor(LazyTensor):
     def _transpose_nonbatch(self):
         return self.__class__(*(lazy_tensor._transpose_nonbatch() for lazy_tensor in self.lazy_tensors), **self._kwargs)
 
-    def _batch_get_indices(self, batch_indices, left_indices, right_indices):
+    def _get_indices(self, left_indices, right_indices, *batch_indices):
         res = torch.ones(left_indices.size(), dtype=self.dtype, device=self.device)
         left_size = self.size(-2)
         right_size = self.size(-1)
@@ -203,22 +203,7 @@ class KroneckerProductLazyTensor(LazyTensor):
             left_indices_i = left_indices.div(left_size)
             right_indices_i = right_indices.div(right_size)
 
-            res = res * lazy_tensor._batch_get_indices(batch_indices, left_indices_i, right_indices_i)
-            left_indices = left_indices - (left_indices_i * left_size)
-            right_indices = right_indices - (right_indices_i * right_size)
-        return res
-
-    def _get_indices(self, left_indices, right_indices):
-        res = torch.ones(left_indices.size(), dtype=self.dtype, device=self.device)
-        left_size = self.size(-2)
-        right_size = self.size(-1)
-        for lazy_tensor in list(self.lazy_tensors)[::-1]:
-            left_size = left_size / lazy_tensor.size(-2)
-            right_size = right_size / lazy_tensor.size(-1)
-            left_indices_i = left_indices.div(left_size)
-            right_indices_i = right_indices.div(right_size)
-
-            res = res * lazy_tensor._get_indices(left_indices_i, right_indices_i)
+            res = res * lazy_tensor._get_indices(left_indices_i, right_indices_i, *batch_indices)
             left_indices = left_indices - (left_indices_i * left_size)
             right_indices = right_indices - (right_indices_i * right_size)
         return res
diff --git a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
index af2f6358..5fa313cd 100644
--- a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
+++ b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
@@ -44,22 +44,11 @@ class LazyEvaluatedKernelTensor(LazyTensor):
     def _transpose_nonbatch(self):
         return self.__class__(self.kernel, self.x2, self.x1, **self.params)
 
-    def _batch_get_indices(self, batch_indices, left_indices, right_indices):
+    def _get_indices(self, left_indices, right_indices, *batch_indices):
         from ..kernels import Kernel
 
-        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(-1, -2), x2.transpose(-1, -2))
-        if isinstance(res, LazyTensor):
-            res = res.evaluate()
-        res = res.view(-1)
-        return res
-
-    def _get_indices(self, left_indices, right_indices):
-        from ..kernels import Kernel
-
-        x1 = self.x1[left_indices, :].unsqueeze(0)
-        x2 = self.x2[right_indices, :].unsqueeze(0)
+        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()
@@ -222,8 +211,24 @@ class LazyEvaluatedKernelTensor(LazyTensor):
         return LazyEvaluatedKernelTensor(self.kernel, x1, x2, **self.params)
 
     def __getitem__(self, index):
-        index = list(index) if isinstance(index, tuple) else [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:
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index c503b565..80f89ffb 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -45,7 +45,7 @@ class LazyTensor(object):
     * :func:`~gpytorch.lazy.LazyTensor._transpose_nonbatch`, which returns a transposed version of the LazyTensor
 
     In addition to these, a LazyTensor may need to define the :func:`~gpytorch.lazy.LazyTensor._transpose_nonbatch`,
-    :func:`~gpytorch.lazy.LazyTensor._get_indices`, and :func:`~gpytorch.lazy.LazyTensor._batch_get_indices`
+    :func:`~gpytorch.lazy.LazyTensor._get_indices`, and :func:`~gpytorch.lazy.LazyTensor._get_indices`
     functions in special cases. See the documentation for these methods for details.
 
     .. note::
@@ -453,7 +453,7 @@ class LazyTensor(object):
             batch_iter = torch.arange(0, size[0], dtype=torch.long, device=self.device)
             batch_iter = batch_iter.unsqueeze(1).repeat(1, size[1]).view(-1)
             row_col_iter = row_col_iter.unsqueeze(1).repeat(size[0], 1).view(-1)
-            return self._batch_get_indices(batch_iter, row_col_iter, row_col_iter).view(size[0], size[1])
+            return self._get_indices(row_col_iter, row_col_iter, batch_iter).view(size[0], size[1])
         else:
             return self._get_indices(row_col_iter, row_col_iter)
 
@@ -533,13 +533,10 @@ class LazyTensor(object):
 
         if precomputed_cache is None:
             train_mean = full_mean.narrow(-1, 0, num_train)
-            if self.ndimension() == 3:
-                if non_batch_train:
-                    train_train_covar = self[0, :num_train, :num_train]
-                else:
-                    train_train_covar = self[:, :num_train, :num_train]
+            if non_batch_train and self.dim() == 3:
+                train_train_covar = self[0, :num_train, :num_train]
             else:
-                train_train_covar = self[:num_train, :num_train]
+                train_train_covar = self[..., :num_train, :num_train]
 
             train_mean = full_mean.narrow(-1, 0, train_train_covar.size(-1))
             if non_batch_train and train_mean.dim() == 2:
@@ -592,14 +589,9 @@ class LazyTensor(object):
         """
         from ..distributions import MultivariateNormal
 
-        if self.ndimension() == 3:
-            train_train_covar = self[:, :num_train, :num_train]
-            test_train_covar = self[:, num_train:, :num_train]
-            test_test_covar = self[:, num_train:, num_train:]
-        else:
-            train_train_covar = self[:num_train, :num_train]
-            test_train_covar = self[num_train:, :num_train]
-            test_test_covar = self[num_train:, num_train:]
+        train_train_covar = self[..., :num_train, :num_train]
+        test_train_covar = self[..., num_train:, :num_train]
+        test_test_covar = self[..., num_train:, num_train:]
 
         train_train_covar = likelihood(
             MultivariateNormal(torch.zeros(1), train_train_covar), train_inputs
@@ -1246,11 +1238,29 @@ class LazyTensor(object):
         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.
         """
+        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]
-        ndimension = self.ndimension()
+
+        # 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))
-        components = list(self._args)
 
         # Normal case if we're indexing the LT with ints or slices
         # Also squeeze dimensions if we're indexing with tensors
@@ -1296,7 +1306,7 @@ class LazyTensor(object):
             return new_lazy_tensor
 
         # Special case: if both row and col are tensor indexed, then we use _get_indices
-        # or _batch_get_indices
+        # or _get_indices
         if torch.is_tensor(left_index) and torch.is_tensor(right_index):
             if left_index.numel() != right_index.numel():
                 raise RuntimeError(
@@ -1317,14 +1327,14 @@ class LazyTensor(object):
                                 batch_index.numel(), left_index.numel(), right_index.numel()
                             )
                         )
-                    return new_lazy_tensor._batch_get_indices(batch_index, left_index, right_index)
+                    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._batch_get_indices(batch_index, left_index, right_index)
+                    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
diff --git a/gpytorch/lazy/matmul_lazy_tensor.py b/gpytorch/lazy/matmul_lazy_tensor.py
index 5d40af17..1137d8e3 100644
--- a/gpytorch/lazy/matmul_lazy_tensor.py
+++ b/gpytorch/lazy/matmul_lazy_tensor.py
@@ -66,31 +66,10 @@ class MatmulLazyTensor(LazyTensor):
     def _transpose_nonbatch(self, *args):
         return self.__class__(self.right_lazy_tensor._transpose_nonbatch(), self.left_lazy_tensor._transpose_nonbatch())
 
-    def _batch_get_indices(self, batch_indices, left_indices, right_indices):
-        n_indices = left_indices.numel()
-        if n_indices > self.size(-1) * self.size(-2) * self.size(-3):
-            return self._evaluated[batch_indices, left_indices, right_indices]
-
-        else:
-            outer_size = batch_indices.size(0)
-            inner_size = self.left_lazy_tensor.size(-1)
-            inner_indices = torch.arange(0, inner_size, dtype=torch.long, device=self.device)
-
-            # Repeat the indices to get all the appropriate terms
-            batch_indices = _outer_repeat(batch_indices, inner_size)
-            left_indices = _outer_repeat(left_indices, inner_size)
-            right_indices = _outer_repeat(right_indices, inner_size)
-            inner_indices = _inner_repeat(inner_indices, outer_size)
-
-            left_vals = self.left_lazy_tensor._batch_get_indices(batch_indices, left_indices, inner_indices)
-            right_vals = self.right_lazy_tensor._batch_get_indices(batch_indices, inner_indices, right_indices)
-
-            return (left_vals.view(-1, inner_size) * right_vals.view(-1, inner_size)).sum(-1)
-
-    def _get_indices(self, left_indices, right_indices):
+    def _get_indices(self, left_indices, right_indices, *batch_indices):
         n_indices = left_indices.numel()
         if n_indices > self.size(-1) * self.size(-2):
-            return self._evaluated[left_indices, right_indices]
+            return self._evaluated.__getitem__((*batch_indices, left_indices, right_indices))
 
         else:
             outer_size = left_indices.size(0)
@@ -98,12 +77,13 @@ class MatmulLazyTensor(LazyTensor):
             inner_indices = torch.arange(0, inner_size, dtype=torch.long, device=self.device)
 
             # Repeat the indices to get all the appropriate terms
+            batch_indices = [_outer_repeat(batch_index, inner_size) for batch_index in batch_indices]
             left_indices = _outer_repeat(left_indices, inner_size)
             right_indices = _outer_repeat(right_indices, inner_size)
             inner_indices = _inner_repeat(inner_indices, outer_size)
 
-            left_vals = self.left_lazy_tensor._get_indices(left_indices, inner_indices)
-            right_vals = self.right_lazy_tensor._get_indices(inner_indices, right_indices)
+            left_vals = self.left_lazy_tensor._get_indices(left_indices, inner_indices, *batch_indices)
+            right_vals = self.right_lazy_tensor._get_indices(inner_indices, right_indices, *batch_indices)
 
             return (left_vals.view(-1, inner_size) * right_vals.view(-1, inner_size)).sum(-1)
 
diff --git a/gpytorch/lazy/mul_lazy_tensor.py b/gpytorch/lazy/mul_lazy_tensor.py
index d1e567cd..acb7258f 100644
--- a/gpytorch/lazy/mul_lazy_tensor.py
+++ b/gpytorch/lazy/mul_lazy_tensor.py
@@ -224,17 +224,10 @@ class MulLazyTensor(LazyTensor):
     def _size(self):
         return self.lazy_tensors[0].size()
 
-    def _batch_get_indices(self, batch_indices, left_indices, right_indices):
-        res = prod(
-            [
-                lazy_tensor._batch_get_indices(batch_indices, left_indices, right_indices)
-                for lazy_tensor in self.lazy_tensors
-            ]
-        )
-        return res
-
-    def _get_indices(self, left_indices, right_indices):
-        res = prod([lazy_tensor._get_indices(left_indices, right_indices) for lazy_tensor in self.lazy_tensors])
+    def _get_indices(self, left_indices, right_indices, *batch_indices):
+        res = prod([
+            lazy_tensor._get_indices(left_indices, right_indices, *batch_indices) for lazy_tensor in self.lazy_tensors
+        ])
         return res
 
     def _transpose_nonbatch(self):
diff --git a/gpytorch/lazy/non_lazy_tensor.py b/gpytorch/lazy/non_lazy_tensor.py
index 813b8b1b..845402b2 100644
--- a/gpytorch/lazy/non_lazy_tensor.py
+++ b/gpytorch/lazy/non_lazy_tensor.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import torch
+from .. import settings
 from ..lazy import LazyTensor
 
 
@@ -34,11 +35,8 @@ class NonLazyTensor(LazyTensor):
     def _transpose_nonbatch(self):
         return NonLazyTensor(self.tensor.transpose(-1, -2))
 
-    def _batch_get_indices(self, batch_indices, left_indices, right_indices):
-        return self.tensor[batch_indices, left_indices, right_indices]
-
-    def _get_indices(self, left_indices, right_indices):
-        return self.tensor[left_indices, right_indices]
+    def _get_indices(self, left_indices, right_indices, *batch_indices):
+        return self.tensor.__getitem__((*batch_indices, left_indices, right_indices))
 
     def diag(self):
         if self.tensor.ndimension() < 3:
@@ -51,10 +49,33 @@ class NonLazyTensor(LazyTensor):
         return self.tensor
 
     def __getitem__(self, index):
-        res = self.tensor[index]
+        # Make sure index is a list
         if not isinstance(index, tuple):
-            index = (index,)
+            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):
@@ -63,5 +84,4 @@ class NonLazyTensor(LazyTensor):
             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/root_lazy_tensor.py b/gpytorch/lazy/root_lazy_tensor.py
index 6d6b3c54..af12449e 100644
--- a/gpytorch/lazy/root_lazy_tensor.py
+++ b/gpytorch/lazy/root_lazy_tensor.py
@@ -54,31 +54,10 @@ class RootLazyTensor(LazyTensor):
     def _transpose_nonbatch(self):
         return self
 
-    def _batch_get_indices(self, batch_indices, left_indices, right_indices):
-        n_indices = left_indices.numel()
-        if n_indices > self.size(-1) * self.size(-2) * self.size(-3):
-            return self.evaluate()[batch_indices, left_indices, right_indices]
-
-        else:
-            outer_size = batch_indices.size(0)
-            inner_size = self.root.size(-1)
-            inner_indices = torch.arange(0, inner_size, dtype=torch.long, device=self.device)
-
-            # Repeat the indices to get all the appropriate terms
-            batch_indices = _outer_repeat(batch_indices, inner_size)
-            left_indices = _outer_repeat(left_indices, inner_size)
-            right_indices = _outer_repeat(right_indices, inner_size)
-            inner_indices = _inner_repeat(inner_indices, outer_size)
-
-            left_vals = self.root._batch_get_indices(batch_indices, left_indices, inner_indices)
-            right_vals = self.root.transpose(-1, -2)._batch_get_indices(batch_indices, inner_indices, right_indices)
-
-            return (left_vals.view(-1, inner_size) * right_vals.view(-1, inner_size)).sum(-1)
-
-    def _get_indices(self, left_indices, right_indices):
+    def _get_indices(self, left_indices, right_indices, *batch_indices):
         n_indices = left_indices.numel()
         if n_indices > self.size(-1) * self.size(-2):
-            return self.evaluate()[left_indices, right_indices]
+            return self.evaluate().__getitem__((*batch_indices, left_indices, right_indices))
 
         else:
             outer_size = left_indices.size(0)
@@ -86,12 +65,13 @@ class RootLazyTensor(LazyTensor):
             inner_indices = torch.arange(0, inner_size, dtype=torch.long, device=self.device)
 
             # Repeat the indices to get all the appropriate terms
+            batch_indices = [_outer_repeat(batch_index, inner_size) for batch_index in batch_indices]
             left_indices = _outer_repeat(left_indices, inner_size)
             right_indices = _outer_repeat(right_indices, inner_size)
             inner_indices = _inner_repeat(inner_indices, outer_size)
 
-            left_vals = self.root._get_indices(left_indices, inner_indices)
-            right_vals = self.root.t()._get_indices(inner_indices, right_indices)
+            left_vals = self.root._get_indices(left_indices, inner_indices, *batch_indices)
+            right_vals = self.root.transpose(-1, -2)._get_indices(inner_indices, right_indices, *batch_indices)
 
         return (left_vals.view(-1, inner_size) * right_vals.view(-1, inner_size)).sum(-1)
 
diff --git a/gpytorch/lazy/sum_batch_lazy_tensor.py b/gpytorch/lazy/sum_batch_lazy_tensor.py
index 095922ba..a4a3e289 100644
--- a/gpytorch/lazy/sum_batch_lazy_tensor.py
+++ b/gpytorch/lazy/sum_batch_lazy_tensor.py
@@ -2,6 +2,7 @@
 
 import torch
 from .block_lazy_tensor import BlockLazyTensor
+from .. import settings
 
 
 class SumBatchLazyTensor(BlockLazyTensor):
@@ -76,26 +77,30 @@ class SumBatchLazyTensor(BlockLazyTensor):
             inner_batch_size = self.base_lazy_tensor.size(0) // self.num_blocks
             return torch.Size([inner_batch_size] + list(base_size)[1:])
 
-    def _batch_get_indices(self, batch_indices, left_indices, right_indices):
-        inner_batch_indices = torch.arange(0, self.num_blocks, dtype=torch.long, device=left_indices.device).unsqueeze_(
-            1
-        )
-        batch_indices = batch_indices.mul(self.num_blocks).unsqueeze_(0).add(inner_batch_indices).view(-1)
-        left_indices = left_indices.unsqueeze(0).repeat(self.num_blocks, 1).view(-1)
-        right_indices = right_indices.unsqueeze(0).repeat(self.num_blocks, 1).view(-1)
-        res = self.base_lazy_tensor._batch_get_indices(batch_indices, left_indices, right_indices)
-
-        res = self.base_lazy_tensor._batch_get_indices(batch_indices, left_indices, right_indices)
-        res = res.view(self.num_blocks, -1).sum(0)
-        return res
-
-    def _get_indices(self, left_indices, right_indices):
-        batch_indices = torch.arange(0, self.base_lazy_tensor.size(0), dtype=torch.long, device=left_indices.device)
-        batch_indices = batch_indices.unsqueeze(1).repeat(1, len(left_indices)).view(-1)
-        left_indices = left_indices.unsqueeze(1).repeat(self.base_lazy_tensor.size(0), 1).view(-1)
-        right_indices = right_indices.unsqueeze(1).repeat(self.base_lazy_tensor.size(0), 1).view(-1)
-        res = self.base_lazy_tensor._batch_get_indices(batch_indices, left_indices, right_indices)
-        return res.view(self.base_lazy_tensor.size(0), -1).sum(0)
+    def _get_indices(self, left_indices, right_indices, *batch_indices):
+        if self.num_blocks is None:
+            if settings.debug.on():
+                assert(len(batch_indices) == 0)
+            batch_indices = torch.arange(0, self.base_lazy_tensor.size(0), dtype=torch.long, device=left_indices.device)
+            batch_indices = batch_indices.unsqueeze(1).repeat(1, len(left_indices)).view(-1)
+            left_indices = left_indices.unsqueeze(1).repeat(self.base_lazy_tensor.size(0), 1).view(-1)
+            right_indices = right_indices.unsqueeze(1).repeat(self.base_lazy_tensor.size(0), 1).view(-1)
+            res = self.base_lazy_tensor._get_indices(left_indices, right_indices, batch_indices)
+            return res.view(self.base_lazy_tensor.size(0), -1).sum(0)
+        else:
+            if settings.debug.on():
+                assert(len(batch_indices) == 1)
+            batch_indices = batch_indices[0]
+            inner_batch_indices = torch.arange(0, self.num_blocks, dtype=torch.long, device=left_indices.device)
+            inner_batch_indices.unsqueeze_(1)
+            batch_indices = batch_indices.mul(self.num_blocks).unsqueeze_(0).add(inner_batch_indices).view(-1)
+            left_indices = left_indices.unsqueeze(0).repeat(self.num_blocks, 1).view(-1)
+            right_indices = right_indices.unsqueeze(0).repeat(self.num_blocks, 1).view(-1)
+            res = self.base_lazy_tensor._get_indices(left_indices, right_indices, batch_indices)
+
+            res = self.base_lazy_tensor._get_indices(left_indices, right_indices, batch_indices)
+            res = res.view(self.num_blocks, -1).sum(0)
+            return res
 
     def _exact_predictive_covar_inv_quad_form_cache(self, train_train_covar_inv_root, test_train_covar):
         if self.num_blocks is None:
diff --git a/gpytorch/lazy/sum_lazy_tensor.py b/gpytorch/lazy/sum_lazy_tensor.py
index 49dc6fcc..918e1993 100644
--- a/gpytorch/lazy/sum_lazy_tensor.py
+++ b/gpytorch/lazy/sum_lazy_tensor.py
@@ -37,15 +37,12 @@ class SumLazyTensor(LazyTensor):
         lazy_tensors_t = [lazy_tensor.transpose(-1, -2) for lazy_tensor in self.lazy_tensors]
         return SumLazyTensor(*lazy_tensors_t)
 
-    def _batch_get_indices(self, batch_indices, left_indices, right_indices):
+    def _get_indices(self, left_indices, right_indices, *batch_indices):
         return sum(
-            lazy_tensor._batch_get_indices(batch_indices, left_indices, right_indices)
+            lazy_tensor._get_indices(left_indices, right_indices, *batch_indices)
             for lazy_tensor in self.lazy_tensors
         )
 
-    def _get_indices(self, left_indices, right_indices):
-        return sum(lazy_tensor._get_indices(left_indices, right_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(
@@ -91,10 +88,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/toeplitz_lazy_tensor.py b/gpytorch/lazy/toeplitz_lazy_tensor.py
index b73a9cf7..91a9c410 100644
--- a/gpytorch/lazy/toeplitz_lazy_tensor.py
+++ b/gpytorch/lazy/toeplitz_lazy_tensor.py
@@ -36,15 +36,10 @@ class ToeplitzLazyTensor(LazyTensor):
     def _transpose_nonbatch(self):
         return ToeplitzLazyTensor(self.column)
 
-    def _batch_get_indices(self, batch_indices, left_indices, right_indices):
+    def _get_indices(self, left_indices, right_indices, *batch_indices):
         n_grid = self.column.size(-1)
         toeplitz_indices = (left_indices - right_indices).fmod(n_grid).abs().long()
-        return self.column[batch_indices, toeplitz_indices]
-
-    def _get_indices(self, left_indices, right_indices):
-        n_grid = self.column.size(-1)
-        toeplitz_indices = (left_indices - right_indices).fmod(n_grid).abs().long()
-        return self.column.index_select(0, toeplitz_indices)
+        return self.column.__getitem__((*batch_indices, toeplitz_indices))
 
     def add_jitter(self, jitter_val=1e-3):
         jitter = torch.zeros_like(self.column)
diff --git a/gpytorch/lazy/zero_lazy_tensor.py b/gpytorch/lazy/zero_lazy_tensor.py
index c9bd4a60..7e26c5fe 100644
--- a/gpytorch/lazy/zero_lazy_tensor.py
+++ b/gpytorch/lazy/zero_lazy_tensor.py
@@ -1,7 +1,8 @@
 #!/usr/bin/env python3
 
 import torch
-from . import LazyTensor
+from .lazy_tensor import LazyTensor
+from .. import settings
 
 
 class ZeroLazyTensor(LazyTensor):
@@ -144,8 +145,24 @@ class ZeroLazyTensor(LazyTensor):
         return self
 
     def __getitem__(self, index):
-        index = list(index) if isinstance(index, tuple) else [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
diff --git a/gpytorch/variational/variational_strategy.py b/gpytorch/variational/variational_strategy.py
index f99e6b4f..8e611031 100644
--- a/gpytorch/variational/variational_strategy.py
+++ b/gpytorch/variational/variational_strategy.py
@@ -87,20 +87,12 @@ class VariationalStrategy(Module):
             full_output = self.model.forward(full_inputs)
             full_mean, full_covar = full_output.mean, full_output.lazy_covariance_matrix
 
-            if full_covar.dim() == 3:
-                test_mean = full_mean[:, n_induc:]
-                induc_mean = full_mean[:, :n_induc]
-                induc_induc_covar = full_covar[:, :n_induc, :n_induc].add_jitter()
-                induc_data_covar = full_covar[:, :n_induc, n_induc:]
-                data_data_covar = full_covar[:, n_induc:, n_induc:]
-                var_dist_mean = variational_dist.mean
-            else:
-                test_mean = full_mean[n_induc:]
-                induc_mean = full_mean[:n_induc]
-                induc_induc_covar = full_covar[:n_induc, :n_induc].add_jitter()
-                induc_data_covar = full_covar[:n_induc, n_induc:]
-                data_data_covar = full_covar[n_induc:, n_induc:]
-                var_dist_mean = variational_dist.mean
+            test_mean = full_mean[..., n_induc:]
+            induc_mean = full_mean[..., :n_induc]
+            induc_induc_covar = full_covar[..., :n_induc, :n_induc].add_jitter()
+            induc_data_covar = full_covar[..., :n_induc, n_induc:]
+            data_data_covar = full_covar[..., n_induc:, n_induc:]
+            var_dist_mean = variational_dist.mean
 
             # Cache the prior distribution, for faster training
             if self.training:


Test patch:
diff --git a/test/lazy/_lazy_tensor_test_case.py b/test/lazy/_lazy_tensor_test_case.py
index 1beb35b5..01f5a2f0 100644
--- a/test/lazy/_lazy_tensor_test_case.py
+++ b/test/lazy/_lazy_tensor_test_case.py
@@ -91,6 +91,18 @@ class RectangularLazyTensorTestCase(object):
         res = lazy_tensor[0:2, :].evaluate()
         actual = evaluated[0:2, :]
         self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+        res = lazy_tensor[..., 0:2].evaluate()
+        actual = evaluated[..., 0:2]
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+        res = lazy_tensor[0:2, ...].evaluate()
+        actual = evaluated[0:2, ...]
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+        res = lazy_tensor[..., 0:2, 2]
+        actual = evaluated[..., 0:2, 2]
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+        res = lazy_tensor[0:2, ..., 2]
+        actual = evaluated[0:2, ..., 2]
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
 
     def test_getitem_tensor_index(self):
         lazy_tensor = self.create_lazy_tensor()
@@ -105,6 +117,15 @@ class RectangularLazyTensorTestCase(object):
         index = (slice(None, None, None), torch.tensor([0, 0, 1, 2]))
         res, actual = lazy_tensor[index], evaluated[index]
         self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+        index = (torch.tensor([0, 0, 1, 2]), Ellipsis)
+        res, actual = lazy_tensor[index], evaluated[index]
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+        index = (Ellipsis, torch.tensor([0, 0, 1, 2]))
+        res, actual = lazy_tensor[index], evaluated[index]
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+        index = (Ellipsis, torch.tensor([0, 0, 1, 2]), torch.tensor([0, 1, 0, 2]))
+        res, actual = lazy_tensor[index], evaluated[index]
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
 
 
 class LazyTensorTestCase(RectangularLazyTensorTestCase):
@@ -307,6 +328,14 @@ class RectangularBatchLazyTensorTestCase(object):
             actual = evaluated.__getitem__((*batch_index, slice(1, None, None), 2))
             self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
 
+        # Ellipsis
+        res = lazy_tensor.__getitem__((Ellipsis, slice(1, None, None), 2))
+        actual = evaluated.__getitem__((Ellipsis, slice(1, None, None), 2))
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+        res = lazy_tensor.__getitem__((slice(1, None, None), Ellipsis, 2))
+        actual = evaluated.__getitem__((slice(1, None, None), Ellipsis, 2))
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+
     def test_getitem_tensor_index(self):
         lazy_tensor = self.create_lazy_tensor()
         evaluated = self.evaluate_lazy_tensor(lazy_tensor)
@@ -327,6 +356,14 @@ class RectangularBatchLazyTensorTestCase(object):
             res, actual = lazy_tensor[index].evaluate(), evaluated[index]
             self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
 
+        # Ellipsis
+        res = lazy_tensor.__getitem__((Ellipsis, torch.tensor([0, 1, 0, 2]), torch.tensor([1, 2, 0, 1])))
+        actual = evaluated.__getitem__((Ellipsis, torch.tensor([0, 1, 0, 2]), torch.tensor([1, 2, 0, 1])))
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+        res = lazy_tensor.__getitem__((torch.tensor([0, 1, 0, 1]), Ellipsis, torch.tensor([1, 2, 0, 1])))
+        actual = evaluated.__getitem__((torch.tensor([0, 1, 0, 1]), Ellipsis, torch.tensor([1, 2, 0, 1])))
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+
 
 class BatchLazyTensorTestCase(RectangularBatchLazyTensorTestCase):
     should_test_sample = False
diff --git a/test/lazy/test_zero_lazy_tensor.py b/test/lazy/test_zero_lazy_tensor.py
index eaca20ce..d2dbec4d 100644
--- a/test/lazy/test_zero_lazy_tensor.py
+++ b/test/lazy/test_zero_lazy_tensor.py
@@ -35,6 +35,17 @@ class TestZeroLazyTensor(unittest.TestCase):
         self.assertLess(torch.norm(res_two - torch.zeros(5, 2, 3)), 1e-4)
         self.assertLess(torch.norm(res_three - torch.zeros(5, 4, 2)), 1e-4)
 
+    def test_getitem_ellipsis(self):
+        lv = ZeroLazyTensor(5, 4, 3)
+
+        res_one = lv[[0, 1]].evaluate()
+        res_two = lv[:, [0, 1], ...].evaluate()
+        res_three = lv[..., [0, 2]].evaluate()
+
+        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)
+        self.assertLess(torch.norm(res_three - torch.zeros(5, 4, 2)), 1e-4)
+
     def test_get_item_tensor_index(self):
         # Tests the default LV.__getitem__ behavior
         lazy_tensor = ZeroLazyTensor(5, 5)
@@ -46,6 +57,10 @@ class TestZeroLazyTensor(unittest.TestCase):
         self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
         index = (slice(None, None, None), torch.tensor([0, 0, 1, 2]))
         self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (Ellipsis, slice(None, None, None), torch.tensor([0, 0, 1, 2]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (Ellipsis, torch.tensor([0, 0, 1, 2]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
 
     def test_get_item_tensor_index_on_batch(self):
         # Tests the default LV.__getitem__ behavior
@@ -68,6 +83,8 @@ class TestZeroLazyTensor(unittest.TestCase):
         self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
         index = (torch.tensor([0, 0, 1, 0]), slice(None, None, None), torch.tensor([0, 0, 1, 1]))
         self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (Ellipsis, torch.tensor([0, 1, 1, 0]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
 
     def test_add_diag(self):
         diag = torch.tensor(1.5)