Repository: cornellius-gp/gpytorch (python)

Issue:
Batch mode InterpolatedLazyVariable __getitem__ doesn't behave like torch.Tensor.__getitem__
`InterpolatedLazyVariable`'s `__getitem__` method has different behavior than `torch.Tensor` in batch mode. Here's a simple example:

```python
matrix = InterpolatedLazyVariable(...) # Some batch mode ILV, say 2 x 10 x 10
eval_matrix = matrix.evaluate() # Get a torch.Tensor of the same matrix.

 # Returns a 2 x 11 tensor containing eval_matrix[0, 0, :] and eval_matrix[1, 2, :]
res1 = eval_matrix[[0, 1], [0, 2], :]

# Returns a 2 x 2 x 11 tensor containing res1 replicated twice.
res2 = matrix[[0, 1], [0, 2], :].evaluate()
```

@gpleiss 

Pull request description:
Fix LazyTensor.__getitem__
This fixes #159, and is a step towards #236.

__getitem__ would fail for batch LazyTensors if you supplied two or more tensor-based indices. E.g. `lazy_tensor[torch.tensor([0, 0, 1]), torch.tensor([2, 1, 2]), :]`.

This fixes the default LT.__getitem__, all custom LT.__getitem__s, and adds a lot of unit tests.

Patch:
diff --git a/gpytorch/lazy/block_diag_lazy_tensor.py b/gpytorch/lazy/block_diag_lazy_tensor.py
index 1e8eadfd..64f51d9f 100644
--- a/gpytorch/lazy/block_diag_lazy_tensor.py
+++ b/gpytorch/lazy/block_diag_lazy_tensor.py
@@ -74,7 +74,7 @@ class BlockDiagLazyTensor(BlockLazyTensor):
         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 = 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)
@@ -85,7 +85,7 @@ class BlockDiagLazyTensor(BlockLazyTensor):
         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 = 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)
diff --git a/gpytorch/lazy/block_lazy_tensor.py b/gpytorch/lazy/block_lazy_tensor.py
index d281a3d1..057c0dee 100644
--- a/gpytorch/lazy/block_lazy_tensor.py
+++ b/gpytorch/lazy/block_lazy_tensor.py
@@ -12,8 +12,7 @@ class BlockLazyTensor(LazyTensor):
     An abstract LazyTensor class for block tensors.
     Super classes will determine how the different blocks are layed out
     (e.g. block diagonal, sum over blocks, etc.)
-
-    BlockLazyTensors represent the groups of blocks as a batched Tensor.
+BlockLazyTensors represent the groups of blocks as a batched Tensor.
     For example, a `k x n x n` tensor represents `k` `n x n` blocks.
 
     For a batched block tensor, the batch dimension is used to represent
@@ -63,9 +62,10 @@ class BlockLazyTensor(LazyTensor):
         # Cases for when there's an inner batch
         else:
             batch_index = index if isinstance(index, int) else index[0]
+            first_tensor_index_dim = None
 
             # Keeping all batch dimensions - recursion base case
-            if batch_index == slice(None, None, None):
+            if isinstance(batch_index, slice) and batch_index == slice(None, None, None):
                 res = super(BlockLazyTensor, self).__getitem__(index)
                 return res
 
@@ -81,6 +81,14 @@ class BlockLazyTensor(LazyTensor):
                 batch_index = slice(start * self.num_blocks, stop * self.num_blocks, step)
                 num_blocks = self.num_blocks
 
+            # Keep sum_batch_index, because we still have an inner batch
+            # Also keep track that there has been tensor indexing
+            elif torch.is_tensor(batch_index):
+                block_index = torch.arange(0, self.num_blocks, dtype=torch.long, device=self.device)
+                batch_index = (batch_index.unsqueeze(1).mul(self.num_blocks) + block_index.unsqueeze(0)).view(-1)
+                num_blocks = self.num_blocks
+                first_tensor_index_dim = 0
+
             else:
                 raise RuntimeError("Unknown batch index type")
 
@@ -94,9 +102,61 @@ class BlockLazyTensor(LazyTensor):
 
             # Else - recurse
             else:
-                if new_var.num_blocks is None:
-                    index = index[1:]
-                else:
-                    index = list(index)
-                    index[0] = slice(None, None, None)
-                return new_var.__getitem__(tuple(index))
+                left_index = index[1]
+                right_index = index[2] if len(index) >= 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
+                squeeze_left = False
+                squeeze_right = False
+                if isinstance(left_index, int):
+                    left_index = slice(left_index, left_index + 1, None)
+                    squeeze_left = True
+                elif torch.is_tensor(left_index):
+                    squeeze_left = True
+                if isinstance(right_index, int):
+                    right_index = slice(right_index, right_index + 1, None)
+                    squeeze_right = True
+                elif torch.is_tensor(right_index):
+                    squeeze_right = True
+
+                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_var.ndimension() == 2:
+                        return new_var._get_indices(left_index, right_index)
+
+                    else:
+                        batch_index = torch.arange(0, new_var.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_var._batch_get_indices(batch_index, left_index, right_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)
+                            return res.view(batch_size, row_col_size)
+
+                # Normal case: we have to do some processing on eithe rthe rows or columns
+                res = new_var._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
diff --git a/gpytorch/lazy/interpolated_lazy_tensor.py b/gpytorch/lazy/interpolated_lazy_tensor.py
index e5f38239..f49e3e0c 100644
--- a/gpytorch/lazy/interpolated_lazy_tensor.py
+++ b/gpytorch/lazy/interpolated_lazy_tensor.py
@@ -619,63 +619,48 @@ class InterpolatedLazyTensor(LazyTensor):
             )
             return res.permute(1, 0).contiguous()
 
-    def __getitem__(self, index):
-        index = list(index) if isinstance(index, tuple) else [index]
+    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.
+        Implementing this is not necessary, but it improves performance
+
+        Args:
+            row_index (slice or LongTensor): index over rows
+            col_index (slice or LongTensor): index over columns
+            first_tensor_index_dim (int or None): first batch dim to have a tensor index (default: None)
+
+        Returns:
+            LazyTensor
+        """
         ndimension = self.ndimension()
-        index += [slice(None, None, None)] * (ndimension - len(index))
-
-        # Check that left interp index and right interp indices are not scalar values
-        squeeze_left = False
-        squeeze_right = False
-        if isinstance(index[-2], int):
-            index[-2] = slice(index[-2], index[-2] + 1, None)
-            squeeze_left = True
-        if isinstance(index[-1], int):
-            index[-1] = slice(index[-1], index[-1] + 1, None)
-            squeeze_right = True
-
-        base_lazy_tensor = self.base_lazy_tensor
+
         left_interp_indices = self.left_interp_indices
         left_interp_values = self.left_interp_values
         right_interp_indices = self.right_interp_indices
         right_interp_values = self.right_interp_values
 
-        # Handle batch dimensions
-        isbatch = ndimension >= 3
-        if isbatch:
-            batch_index = tuple(index[:-2])
-            base_lazy_tensor = self.base_lazy_tensor[batch_index]
-            left_interp_indices = self.left_interp_indices[batch_index]
-            left_interp_values = self.left_interp_values[batch_index]
-            right_interp_indices = self.right_interp_indices[batch_index]
-            right_interp_values = self.right_interp_values[batch_index]
-
-        ndimension = base_lazy_tensor.ndimension()
+        batch_iter = [slice(None, None, None)] * (ndimension - 2)
+        if first_tensor_index_dim is not None:
+            batch_iter[first_tensor_index_dim] = torch.arange(
+                0, self.size(first_tensor_index_dim), dtype=torch.long, device=self.device
+            )
 
-        # Handle left interp
-        left_index = tuple([slice(None, None, None)] * (ndimension - 2) + [index[-2]])
+        left_index = (*batch_iter, row_index)
         left_interp_indices = left_interp_indices[left_index]
         left_interp_values = left_interp_values[left_index]
+        if first_tensor_index_dim is not None and torch.is_tensor(row_index):
+            left_interp_indices = left_interp_indices.unsqueeze(-2)
+            left_interp_values = left_interp_values.unsqueeze(-2)
 
-        # Handle right interp
-        right_index = tuple([slice(None, None, None)] * (ndimension - 2) + [index[-1]])
+        right_index = (*batch_iter, col_index)
         right_interp_indices = right_interp_indices[right_index]
         right_interp_values = right_interp_values[right_index]
+        if first_tensor_index_dim is not None and torch.is_tensor(col_index):
+            right_interp_indices = right_interp_indices.unsqueeze(-2)
+            right_interp_values = right_interp_values.unsqueeze(-2)
 
         res = self.__class__(
-            base_lazy_tensor,
-            left_interp_indices,
-            left_interp_values,
-            right_interp_indices,
-            right_interp_values,
-            **self._kwargs
+            self.base_lazy_tensor, left_interp_indices, left_interp_values,
+            right_interp_indices, right_interp_values, **self._kwargs
         )
-
-        if squeeze_left or squeeze_right:
-            res = res.evaluate()
-            if squeeze_left:
-                res = res.squeeze(-2)
-            if squeeze_right:
-                res = res.squeeze(-1)
-
         return res
diff --git a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
index 7f7ca56a..7c2f0a63 100644
--- a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
+++ b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
@@ -158,6 +158,7 @@ class LazyEvaluatedKernelTensor(LazyTensor):
         index = list(index) if isinstance(index, tuple) else [index]
         ndimension = self.ndimension()
         index += [slice(None, None, None)] * (ndimension - len(index))
+
         if self.is_batch:
             batch_index = index[0]
             left_index = index[1]
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index 0036681e..50eed1f8 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -96,25 +96,26 @@ class LazyTensor(object):
         """
         return self.diag()
 
-    def _getitem_nonbatch(self, row_index, col_index):
+    def _getitem_nonbatch(self, row_index, col_index, first_tensor_index_dim=None):
         """
-        Given an index over rows and columns,
-        Returns a LazyTensor with those rows and columns selected
-
+        Given an index over rows and columns, gets those items from the LazyTensor.
         Implementing this is not necessary, but it improves performance
 
         Args:
             row_index (slice or LongTensor): index over rows
             col_index (slice or LongTensor): index over columns
+            first_tensor_index_dim (int or None): first batch dim to have a tensor index (default: None)
+
+        Returns:
+            LazyTensor
         """
         from .interpolated_lazy_tensor import InterpolatedLazyTensor
 
         ndimension = self.ndimension()
-
         batch_sizes = list(self.size()[:-2])
+
         left_row_iter = torch.arange(0, self.size()[-2], dtype=torch.long, device=self.device)
         right_row_iter = torch.arange(0, self.size()[-1], dtype=torch.long, device=self.device)
-
         left_interp_indices = left_row_iter[row_index].unsqueeze(-1)
         right_interp_indices = right_row_iter[col_index].unsqueeze(-1)
 
@@ -124,9 +125,19 @@ class LazyTensor(object):
             left_interp_indices.unsqueeze_(0)
             right_interp_indices.unsqueeze_(0)
 
-        left_interp_indices = left_interp_indices.expand(*(batch_sizes + [left_interp_len, 1]))
+        if first_tensor_index_dim is not None and torch.is_tensor(row_index):
+            view_size = [1] * ndimension
+            view_size[first_tensor_index_dim] = left_interp_indices.numel()
+            left_interp_indices = left_interp_indices.view(*view_size).expand(*(batch_sizes + [1, 1]))
+        else:
+            left_interp_indices = left_interp_indices.expand(*(batch_sizes + [left_interp_len, 1]))
         left_interp_values = torch.ones(left_interp_indices.size(), dtype=self.dtype, device=self.device)
-        right_interp_indices = right_interp_indices.expand(*(batch_sizes + [right_interp_len, 1]))
+        if first_tensor_index_dim is not None and torch.is_tensor(col_index):
+            view_size = [1] * ndimension
+            view_size[first_tensor_index_dim] = right_interp_indices.numel()
+            right_interp_indices = right_interp_indices.view(*view_size).expand(*(batch_sizes + [1, 1]))
+        else:
+            right_interp_indices = right_interp_indices.expand(*(batch_sizes + [right_interp_len, 1]))
         right_interp_values = torch.ones(right_interp_indices.size(), dtype=self.dtype, device=self.device)
 
         res = InterpolatedLazyTensor(
@@ -1060,28 +1071,41 @@ class LazyTensor(object):
         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
         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]
 
+            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)
@@ -1090,13 +1114,46 @@ class LazyTensor(object):
         ):
             return new_lazy_tensor
 
-        res = new_lazy_tensor._getitem_nonbatch(left_index, right_index)
-        if squeeze_left or squeeze_right:
+        # Special case: if both row and col are tensor indexed, then we use _get_indices
+        # or _batch_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._batch_get_indices(batch_index, left_index, right_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)
+                    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)
+        if squeeze_left:
+            res = res.squeeze(-2)
+        if squeeze_right:
+            res = res.squeeze(-1)
 
         return res
 
diff --git a/gpytorch/lazy/zero_lazy_tensor.py b/gpytorch/lazy/zero_lazy_tensor.py
index 8cb92080..5c2a43a0 100644
--- a/gpytorch/lazy/zero_lazy_tensor.py
+++ b/gpytorch/lazy/zero_lazy_tensor.py
@@ -12,10 +12,21 @@ class ZeroLazyTensor(LazyTensor):
     Special LazyTensor representing zero.
     """
 
-    def __init__(self, *sizes):
+    def __init__(self, *sizes, dtype=None, device=None):
         super(ZeroLazyTensor, self).__init__(*sizes)
         self.sizes = list(sizes)
 
+        self._dtype = dtype or torch.get_default_dtype()
+        self._device = device or torch.device("cpu")
+
+    @property
+    def dtype(self):
+        return self._dtype
+
+    @property
+    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):
@@ -139,13 +150,29 @@ class ZeroLazyTensor(LazyTensor):
         index = list(index) if isinstance(index, tuple) else [index]
         ndimension = self.ndimension()
         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 += [len(torch.arange(self.size(ix))[sub_index])]
+                new_sizes.append(len(torch.arange(self.size(ix))[sub_index]))
             else:
-                new_sizes += [len(sub_index)]
+                new_sizes.append(len(sub_index))
 
-        return ZeroLazyTensor(*new_sizes)
+        if evaluate:
+            return torch.zeros(*new_sizes, dtype=self.dtype, device=self.device)
+        else:
+            return ZeroLazyTensor(*new_sizes)
diff --git a/gpytorch/utils/__init__.py b/gpytorch/utils/__init__.py
index 2fea4872..48b08a21 100644
--- a/gpytorch/utils/__init__.py
+++ b/gpytorch/utils/__init__.py
@@ -177,10 +177,13 @@ def left_t_interp(interp_indices, interp_values, rhs, output_dim):
 def prod(items):
     """
     """
-    res = items[0]
-    for item in items[1:]:
-        res = res * item
-    return res
+    if len(items):
+        res = items[0]
+        for item in items[1:]:
+            res = res * item
+        return res
+    else:
+        return 1
 
 
 def sparse_eye(size):


Test patch:
diff --git a/test/lazy/test_block_diag_lazy_tensor.py b/test/lazy/test_block_diag_lazy_tensor.py
index a5c3599b..28b2cac7 100644
--- a/test/lazy/test_block_diag_lazy_tensor.py
+++ b/test/lazy/test_block_diag_lazy_tensor.py
@@ -121,6 +121,42 @@ class TestBlockDiagLazyTensor(unittest.TestCase):
         actual = actual_block_diag[1:, :5, 2]
         self.assertTrue(approx_equal(actual, res))
 
+    def test_get_item_tensor_index(self):
+        # Tests the default LV.__getitem__ behavior
+        block_tensor = self.blocks.clone().requires_grad_(True)
+        lazy_tensor = BlockDiagLazyTensor(NonLazyTensor(block_tensor))
+        evaluated = lazy_tensor.evaluate()
+
+        index = (torch.tensor([0, 0, 1, 2]), torch.tensor([0, 1, 0, 2]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 0, 1, 2]), slice(None, None, None))
+        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]))
+
+    def test_get_item_tensor_index_on_batch(self):
+        # Tests the default LV.__getitem__ behavior
+        block_tensor = self.blocks.clone().requires_grad_(True)
+        lazy_tensor = BlockDiagLazyTensor(NonLazyTensor(block_tensor), num_blocks=4)
+        evaluated = lazy_tensor.evaluate()
+
+        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), torch.tensor([1, 2, 0, 1]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), slice(None, None, None))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 1, 1]), slice(None, None, None), torch.tensor([0, 1, 2]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (slice(None, None, None), torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 0, 1, 1]), slice(None, None, None), slice(None, None, None))
+        self.assertTrue(approx_equal(lazy_tensor[index].evaluate(), evaluated[index]))
+        index = (slice(None, None, None), torch.tensor([0, 0, 1, 2]), torch.tensor([0, 0, 1, 1]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), slice(None, None, None))
+        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]))
+
     def test_sample(self):
         block_tensor = self.blocks.clone().requires_grad_(True)
         res = BlockDiagLazyTensor(NonLazyTensor(block_tensor))
diff --git a/test/lazy/test_interpolated_lazy_tensor.py b/test/lazy/test_interpolated_lazy_tensor.py
index ca5de0cd..e5740658 100644
--- a/test/lazy/test_interpolated_lazy_tensor.py
+++ b/test/lazy/test_interpolated_lazy_tensor.py
@@ -299,6 +299,63 @@ class TestInterpolatedLazyTensor(unittest.TestCase):
         self.assertTrue(approx_equal(interp_lazy_tensor[1, 1, :2], actual[1, 1, :2]))
         self.assertTrue(approx_equal(interp_lazy_tensor[1, :1, 2], actual[1, :1, 2]))
 
+    def test_get_item_tensor_index(self):
+        # Tests the default LV.__getitem__ behavior
+        left_interp_indices = torch.LongTensor([[2, 3], [3, 4], [4, 5]])
+        left_interp_values = torch.tensor([[1, 1], [1, 1], [1, 1]], dtype=torch.float)
+        right_interp_indices = torch.LongTensor([[0, 1], [1, 2], [2, 3]])
+        right_interp_values = torch.tensor([[1, 1], [1, 1], [1, 1]], dtype=torch.float)
+
+        base_lazy_tensor_mat = torch.randn(6, 6)
+        base_lazy_tensor_mat = base_lazy_tensor_mat.transpose(-1, -2).matmul(base_lazy_tensor_mat)
+
+        base_lazy_tensor = NonLazyTensor(base_lazy_tensor_mat)
+        interp_lazy_tensor = InterpolatedLazyTensor(
+            base_lazy_tensor, left_interp_indices, left_interp_values, right_interp_indices, right_interp_values
+        )
+        evaluated = interp_lazy_tensor.evaluate()
+
+        index = (torch.tensor([0, 0, 1, 2]), torch.tensor([0, 1, 0, 2]))
+        self.assertTrue(approx_equal(interp_lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 0, 1, 2]), slice(None, None, None))
+        self.assertTrue(approx_equal(interp_lazy_tensor[index], evaluated[index]))
+        index = (slice(None, None, None), torch.tensor([0, 0, 1, 2]))
+        self.assertTrue(approx_equal(interp_lazy_tensor[index], evaluated[index]))
+
+    def test_get_item_tensor_index_on_batch(self):
+        # Tests the default LV.__getitem__ behavior
+        left_interp_indices = torch.LongTensor([[2, 3], [3, 4], [4, 5]]).repeat(5, 1, 1)
+        left_interp_values = torch.tensor([[1, 1], [1, 1], [1, 1]], dtype=torch.float).repeat(5, 1, 1)
+        right_interp_indices = torch.LongTensor([[0, 1], [1, 2], [2, 3]]).repeat(5, 1, 1)
+        right_interp_values = torch.tensor([[1, 1], [1, 1], [1, 1]], dtype=torch.float).repeat(5, 1, 1)
+
+        base_lazy_tensor_mat = torch.randn(5, 6, 6)
+        base_lazy_tensor_mat = base_lazy_tensor_mat.transpose(1, 2).matmul(base_lazy_tensor_mat)
+        base_lazy_tensor_mat.requires_grad = True
+
+        base_lazy_tensor = NonLazyTensor(base_lazy_tensor_mat)
+        interp_lazy_tensor = InterpolatedLazyTensor(
+            base_lazy_tensor, left_interp_indices, left_interp_values, right_interp_indices, right_interp_values
+        )
+        evaluated = interp_lazy_tensor.evaluate()
+
+        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), torch.tensor([1, 2, 0, 1]))
+        self.assertTrue(approx_equal(interp_lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), slice(None, None, None))
+        self.assertTrue(approx_equal(interp_lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 1, 1]), slice(None, None, None), torch.tensor([0, 1, 2]))
+        self.assertTrue(approx_equal(interp_lazy_tensor[index], evaluated[index]))
+        index = (slice(None, None, None), torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]))
+        self.assertTrue(approx_equal(interp_lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 0, 1, 1]), slice(None, None, None), slice(None, None, None))
+        self.assertTrue(approx_equal(interp_lazy_tensor[index].evaluate(), evaluated[index]))
+        index = (slice(None, None, None), torch.tensor([0, 0, 1, 2]), torch.tensor([0, 0, 1, 1]))
+        self.assertTrue(approx_equal(interp_lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), slice(None, None, None))
+        self.assertTrue(approx_equal(interp_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(interp_lazy_tensor[index], evaluated[index]))
+
     def test_diag(self):
         left_interp_indices = torch.LongTensor([[2, 3], [3, 4], [4, 5]])
         left_interp_values = torch.tensor([[1, 1], [1, 1], [1, 1]], dtype=torch.float)
diff --git a/test/lazy/test_sum_batch_lazy_tensor.py b/test/lazy/test_sum_batch_lazy_tensor.py
index c0f34a34..f00bd43e 100644
--- a/test/lazy/test_sum_batch_lazy_tensor.py
+++ b/test/lazy/test_sum_batch_lazy_tensor.py
@@ -110,6 +110,42 @@ class TestSumBatchLazyTensor(unittest.TestCase):
         actual = actual_mat[1:, :5, 2]
         self.assertTrue(approx_equal(actual, res))
 
+    def test_get_item_tensor_index(self):
+        # Tests the default LV.__getitem__ behavior
+        block_tensor = self.blocks.clone().requires_grad_(True)
+        lazy_tensor = SumBatchLazyTensor(NonLazyTensor(block_tensor))
+        evaluated = lazy_tensor.evaluate()
+
+        index = (torch.tensor([0, 0, 1, 2]), torch.tensor([0, 1, 0, 2]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 0, 1, 2]), slice(None, None, None))
+        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]))
+
+    def test_get_item_tensor_index_on_batch(self):
+        # Tests the default LV.__getitem__ behavior
+        block_tensor = self.blocks.clone().requires_grad_(True)
+        lazy_tensor = SumBatchLazyTensor(NonLazyTensor(block_tensor), num_blocks=4)
+        evaluated = lazy_tensor.evaluate()
+
+        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), torch.tensor([1, 2, 0, 1]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), slice(None, None, None))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 1, 1]), slice(None, None, None), torch.tensor([0, 1, 2]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (slice(None, None, None), torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 0, 1, 1]), slice(None, None, None), slice(None, None, None))
+        self.assertTrue(approx_equal(lazy_tensor[index].evaluate(), evaluated[index]))
+        index = (slice(None, None, None), torch.tensor([0, 0, 1, 2]), torch.tensor([0, 0, 1, 1]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), slice(None, None, None))
+        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]))
+
     def test_sample(self):
         block_tensor = self.blocks.clone().requires_grad_(True)
         res = SumBatchLazyTensor(NonLazyTensor(block_tensor))
diff --git a/test/lazy/test_toeplitz_lazy_tensor.py b/test/lazy/test_toeplitz_lazy_tensor.py
index 3a05406e..bb072e94 100644
--- a/test/lazy/test_toeplitz_lazy_tensor.py
+++ b/test/lazy/test_toeplitz_lazy_tensor.py
@@ -53,21 +53,55 @@ class TestToeplitzLazyTensor(unittest.TestCase):
         self.assertTrue(utils.approx_equal(res, actual))
 
     def test_get_item_square_on_tensor(self):
+        # Tests the default LV.__getitem__ behavior
         toeplitz_var = ToeplitzLazyTensor(torch.tensor([1, 2, 3, 4], dtype=torch.float))
         evaluated = toeplitz_var.evaluate()
-
         self.assertTrue(utils.approx_equal(toeplitz_var[2:4, 2:4].evaluate(), evaluated[2:4, 2:4]))
 
+    def test_get_item_tensor_index(self):
+        # Tests the default LV.__getitem__ behavior
+        toeplitz_var = ToeplitzLazyTensor(torch.tensor([1, 2, 3, 4], dtype=torch.float))
+        evaluated = toeplitz_var.evaluate()
+        index = (torch.tensor([0, 0, 1, 2]), torch.tensor([0, 1, 0, 2]))
+        self.assertTrue(utils.approx_equal(toeplitz_var[index], evaluated[index]))
+        index = (torch.tensor([0, 0, 1, 2]), slice(None, None, None))
+        self.assertTrue(utils.approx_equal(toeplitz_var[index], evaluated[index]))
+        index = (slice(None, None, None), torch.tensor([0, 0, 1, 2]))
+        self.assertTrue(utils.approx_equal(toeplitz_var[index], evaluated[index]))
+
     def test_get_item_on_batch(self):
+        # Tests the default LV.__getitem__ behavior
         toeplitz_var = ToeplitzLazyTensor(self.batch_toeplitz_column)
         evaluated = toeplitz_var.evaluate()
         self.assertTrue(utils.approx_equal(toeplitz_var[0, 1:3].evaluate(), evaluated[0, 1:3]))
 
     def test_get_item_scalar_on_batch(self):
+        # Tests the default LV.__getitem__ behavior
         toeplitz_var = ToeplitzLazyTensor(torch.tensor([[1, 2, 3, 4]], dtype=torch.float))
         evaluated = toeplitz_var.evaluate()
         self.assertTrue(utils.approx_equal(toeplitz_var[0].evaluate(), evaluated[0]))
 
+    def test_get_item_tensor_index_on_batch(self):
+        # Tests the default LV.__getitem__ behavior
+        toeplitz_var = ToeplitzLazyTensor(torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=torch.float))
+        evaluated = toeplitz_var.evaluate()
+        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), torch.tensor([1, 2, 0, 3]))
+        self.assertTrue(utils.approx_equal(toeplitz_var[index], evaluated[index]))
+        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), slice(None, None, None))
+        self.assertTrue(utils.approx_equal(toeplitz_var[index], evaluated[index]))
+        index = (torch.tensor([0, 1, 1]), slice(None, None, None), torch.tensor([0, 1, 2]))
+        self.assertTrue(utils.approx_equal(toeplitz_var[index], evaluated[index]))
+        index = (slice(None, None, None), torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]))
+        self.assertTrue(utils.approx_equal(toeplitz_var[index], evaluated[index]))
+        index = (torch.tensor([0, 0, 1, 1]), slice(None, None, None), slice(None, None, None))
+        self.assertTrue(utils.approx_equal(toeplitz_var[index].evaluate(), evaluated[index]))
+        index = (slice(None, None, None), torch.tensor([0, 0, 1, 2]), torch.tensor([0, 0, 1, 1]))
+        self.assertTrue(utils.approx_equal(toeplitz_var[index], evaluated[index]))
+        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), slice(None, None, None))
+        self.assertTrue(utils.approx_equal(toeplitz_var[index], evaluated[index]))
+        index = (torch.tensor([0, 0, 1, 0]), slice(None, None, None), torch.tensor([0, 0, 1, 1]))
+        self.assertTrue(utils.approx_equal(toeplitz_var[index], evaluated[index]))
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/lazy/test_zero_lazy_tensor.py b/test/lazy/test_zero_lazy_tensor.py
index 3fdc77a3..1334aca8 100644
--- a/test/lazy/test_zero_lazy_tensor.py
+++ b/test/lazy/test_zero_lazy_tensor.py
@@ -20,8 +20,8 @@ class TestZeroLazyTensor(unittest.TestCase):
         lv = ZeroLazyTensor(5, 4, 3)
 
         res_one = lv[0].evaluate()
-        res_two = lv[:, 1, :].evaluate()
-        res_three = lv[:, :, 2].evaluate()
+        res_two = lv[:, 1, :]
+        res_three = lv[:, :, 2]
 
         self.assertLess(torch.norm(res_one - torch.zeros(4, 3)), 1e-4)
         self.assertLess(torch.norm(res_two - torch.zeros(5, 3)), 1e-4)
@@ -38,6 +38,40 @@ 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_get_item_tensor_index(self):
+        # Tests the default LV.__getitem__ behavior
+        lazy_tensor = ZeroLazyTensor(5, 5)
+        evaluated = lazy_tensor.evaluate()
+
+        index = (torch.tensor([0, 0, 1, 2]), torch.tensor([0, 1, 0, 2]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 0, 1, 2]), slice(None, None, None))
+        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]))
+
+    def test_get_item_tensor_index_on_batch(self):
+        # Tests the default LV.__getitem__ behavior
+        lazy_tensor = ZeroLazyTensor(3, 5, 5)
+        evaluated = lazy_tensor.evaluate()
+
+        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), torch.tensor([1, 2, 0, 1]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), slice(None, None, None))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 1, 1]), slice(None, None, None), torch.tensor([0, 1, 2]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (slice(None, None, None), torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 0, 1, 1]), slice(None, None, None), slice(None, None, None))
+        self.assertTrue(approx_equal(lazy_tensor[index].evaluate(), evaluated[index]))
+        index = (slice(None, None, None), torch.tensor([0, 0, 1, 2]), torch.tensor([0, 0, 1, 1]))
+        self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
+        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), slice(None, None, None))
+        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]))
+
     def test_add_diag(self):
         diag = torch.tensor(1.5)
         res = ZeroLazyTensor(5, 5).add_diag(diag).evaluate()