Fix slicing of BlockDiagLinearOperator

linear_operator/operators/block_interleaved_linear_operator.py

@@ -9,6 +9,7 @@
 from linear_operator.operators._linear_operator import IndexType, LinearOperator
 from linear_operator.operators.block_linear_operator import BlockLinearOperator
 
+from linear_operator.utils.getitem import _is_noop_index, _noop_index
 from linear_operator.utils.memoize import cached
 
 
@@ -80,6 +81,29 @@ def _get_indices(self, row_index: IndexType, col_index: IndexType, *batch_indice
         res = res * torch.eq(row_index_block, col_index_block).type_as(res)
         return res
 
+    def _getitem(self, row_index: IndexType, col_index: IndexType, *batch_indices: IndexType) -> LinearOperator:
+        # Batch indexing, tensor indices and steps are handled by the base case
+        if (_is_noop_index(row_index) and _is_noop_index(col_index)) or not (
+            isinstance(row_index, slice) and isinstance(col_index, slice)
+        ):
+            return super()._getitem(row_index, col_index, *batch_indices)
+        num_blocks = self.num_blocks
+        num_rows, num_cols = self.matrix_shape
+        row_start, row_end, row_step = row_index.start or 0, row_index.stop or num_rows, row_index.step
+        col_start, col_end, col_step = col_index.start or 0, col_index.stop or num_cols, col_index.step
+        if row_step is not None or col_step is not None:
+            return super()._getitem(row_index, col_index, *batch_indices)
+
+        # Row (column) i of the matrix is row (column) i // num_blocks of block i % num_blocks.
+        # So if the slice bounds are multiples of num_blocks, the slices select the same rows and
+        # columns of every block. Otherwise, use the base case
+        if (row_start % num_blocks) or (col_start % num_blocks) or (row_end % num_blocks) or (col_end % num_blocks):
+            return super()._getitem(row_index, col_index, *batch_indices)
+        row_index = slice(row_start // num_blocks, row_end // num_blocks, None)
+        col_index = slice(col_start // num_blocks, col_end // num_blocks, None)
+        new_base_linear_op = self.base_linear_op._getitem(row_index, col_index, *batch_indices, _noop_index)
+        return self.__class__(new_base_linear_op, block_dim=-3)
+
     def _remove_batch_dim(self, other):
         other = other.transpose(-2, -3).contiguous()
         shape = list(other.shape)

linear_operator/operators/block_linear_operator.py

@@ -70,36 +70,9 @@ def _getitem(self, row_index: IndexType, col_index: IndexType, *batch_indices: I
         if _is_noop_index(row_index) and _is_noop_index(col_index):
             return self.__class__(self.base_linear_op._getitem(row_index, col_index, *batch_indices, _noop_index))
 
-        # If either of the dimensions are indices, it's too complicated - go with the base case
-        if not isinstance(row_index, slice) or not isinstance(col_index, slice):
-            # It's too complicated to deal with tensor indices in this case - we'll use the super method
-            return super()._getitem(row_index, col_index, *batch_indices)
-
-        # Now we know that row_index and col_index
-        num_blocks = self.num_blocks
-        num_rows, num_cols = self.matrix_shape
-        row_start, row_end, row_step = row_index.start or 0, row_index.stop or num_rows, row_index.step
-        col_start, col_end, col_step = col_index.start or 0, col_index.stop or num_cols, col_index.step
-
-        # If we have a step, it's too complicated - go with the base case
-        if row_step is not None or col_step is not None:
-            return super()._getitem(row_index, col_index, *batch_indices)
-
-        # Let's make sure that the slice dimensions perfectly correspond with the number of
-        # outputs per input that we have
-        # Otherwise - its too complicated. We'll go with the base case
-        if (row_start % num_blocks) or (col_start % num_blocks) or (row_end % num_blocks) or (col_end % num_blocks):
-            return super()._getitem(row_index, col_index, *batch_indices)
-
-        # Otherwise - let's divide the slices by the number of outputs per input
-        row_index = slice(row_start // num_blocks, row_end // num_blocks, None)
-        col_index = slice(col_start // num_blocks, col_end // num_blocks, None)
-
-        # Now we can try the super call!
-        new_base_linear_op = self.base_linear_op._getitem(row_index, col_index, *batch_indices, _noop_index)
-
-        # Now construct a kernel with those indices
-        return self.__class__(new_base_linear_op, block_dim=-3)
+        # Which rows and columns of the blocks a slice selects depends on how the blocks are laid out,
+        # so subclasses handle the slices that they can map onto the blocks. Otherwise, use the base case.
+        return super()._getitem(row_index, col_index, *batch_indices)
 
     def _matmul(
         self: LinearOperator,  # shape: (*batch, M, N)

test/operators/test_block_diag_linear_operator.py

@@ -1,14 +1,60 @@
 #!/usr/bin/env python3
 
 import unittest
+from itertools import product
 
 import torch
 
-from linear_operator.operators import BlockDiagLinearOperator, DenseLinearOperator, DiagLinearOperator
+from linear_operator.operators import (
+    BlockDiagLinearOperator,
+    CatLinearOperator,
+    DenseLinearOperator,
+    DiagLinearOperator,
+)
+from linear_operator.test.base_test_case import BaseTestCase
 from linear_operator.test.linear_operator_test_case import LinearOperatorTestCase
 
 
-class TestBlockDiagLinearOperator(LinearOperatorTestCase, unittest.TestCase):
+class _BlockSliceTestMixin:
+    def test_getitem_slices(self):
+        linear_op = self.create_linear_op()
+        evaluated = self.evaluate_linear_op(linear_op)
+        num_blocks, block_size = linear_op.num_blocks, linear_op.base_linear_op.size(-1)
+        size = num_blocks * block_size
+        slices = [
+            slice(0, block_size),  # one block
+            slice(block_size, 3 * block_size),  # two blocks
+            slice(block_size, size),
+            slice(0, num_blocks),  # bounds are multiples of num_blocks
+            slice(num_blocks, 2 * num_blocks),
+            slice(1, num_blocks + 1),
+            slice(-num_blocks, None),
+            slice(None, num_blocks),
+            slice(0, num_blocks, 1),
+            slice(0, 2 * num_blocks, 2),
+        ]
+        for row_index, col_index in product(slices, repeat=2):
+            with self.subTest(row_index=row_index, col_index=col_index):
+                res = linear_op[..., row_index, col_index]
+                self.assertAllClose(res.to_dense(), evaluated[..., row_index, col_index])
+                if linear_op.dim() > 2:
+                    res = linear_op[1, ..., row_index, col_index]
+                    self.assertAllClose(res.to_dense(), evaluated[1, ..., row_index, col_index])
+
+
+class TestBlockDiagLinearOperatorCatBase(_BlockSliceTestMixin, BaseTestCase, unittest.TestCase):
+    # MultitaskMultivariateNormal.from_independent_mvns in GPyTorch concatenates the blocks
+    seed = 0
+
+    def create_linear_op(self):
+        blocks = [DenseLinearOperator(torch.randn(1, 4, 4)) for _ in range(3)]
+        return BlockDiagLinearOperator(CatLinearOperator(*blocks, dim=0))
+
+    def evaluate_linear_op(self, linear_op):
+        return torch.block_diag(*(block.tensor[0] for block in linear_op.base_linear_op.linear_ops))
+
+
+class TestBlockDiagLinearOperator(_BlockSliceTestMixin, LinearOperatorTestCase, unittest.TestCase):
     seed = 0
     should_test_sample = True
 
@@ -38,7 +84,7 @@ class TestBlockDiagLinearOperatorFromTensor(TestBlockDiagLinearOperator):
     _initialize_from_tensor = True
 
 
-class TestBlockDiagLinearOperatorBatch(LinearOperatorTestCase, unittest.TestCase):
+class TestBlockDiagLinearOperatorBatch(_BlockSliceTestMixin, LinearOperatorTestCase, unittest.TestCase):
     seed = 0
     should_test_sample = True
 
@@ -57,7 +103,7 @@ def evaluate_linear_op(self, linear_op):
         return actual
 
 
-class TestBlockDiagLinearOperatorMultiBatch(LinearOperatorTestCase, unittest.TestCase):
+class TestBlockDiagLinearOperatorMultiBatch(_BlockSliceTestMixin, LinearOperatorTestCase, unittest.TestCase):
     seed = 0
     # Because these LTs are large, we'll skil the big tests
     should_test_sample = False

test/operators/test_block_interleaved_linear_operator.py

@@ -1,14 +1,47 @@
 #!/usr/bin/env python3
 
 import unittest
+from itertools import product
 
 import torch
 
 from linear_operator.operators import BlockInterleavedLinearOperator, DenseLinearOperator
 from linear_operator.test.linear_operator_test_case import LinearOperatorTestCase
 
 
-class TestBlockInterleavedLinearOperator(LinearOperatorTestCase, unittest.TestCase):
+class _BlockSliceTestMixin:
+    def test_getitem_slices(self):
+        linear_op = self.create_linear_op()
+        evaluated = self.evaluate_linear_op(linear_op)
+        num_blocks, block_size = linear_op.num_blocks, linear_op.base_linear_op.size(-1)
+        size = num_blocks * block_size
+        slices = [
+            slice(0, num_blocks),  # bounds are multiples of num_blocks
+            slice(num_blocks, 3 * num_blocks),
+            slice(num_blocks, size),
+            slice(0, block_size),  # bounds are multiples of the block size
+            slice(block_size, 2 * block_size),
+            slice(1, num_blocks + 1),
+            slice(-num_blocks, None),
+            slice(None, num_blocks),
+            slice(0, num_blocks, 1),
+            slice(0, 2 * num_blocks, 2),
+        ]
+        for row_index, col_index in product(slices, repeat=2):
+            with self.subTest(row_index=row_index, col_index=col_index):
+                res = linear_op[..., row_index, col_index]
+                self.assertAllClose(res.to_dense(), evaluated[..., row_index, col_index])
+                if linear_op.dim() > 2:
+                    res = linear_op[1, ..., row_index, col_index]
+                    self.assertAllClose(res.to_dense(), evaluated[1, ..., row_index, col_index])
+
+        # Slices whose bounds are multiples of num_blocks keep the interleaved structure
+        res = linear_op[..., num_blocks : 3 * num_blocks, 0 : 2 * num_blocks]
+        self.assertIsInstance(res, BlockInterleavedLinearOperator)
+        self.assertEqual(res.num_blocks, num_blocks)
+
+
+class TestBlockInterleavedLinearOperator(_BlockSliceTestMixin, LinearOperatorTestCase, unittest.TestCase):
     seed = 0
     should_test_sample = True
 
@@ -28,7 +61,7 @@ def evaluate_linear_op(self, linear_op):
         return actual
 
 
-class TestBlockInterleavedLinearOperatorBatch(LinearOperatorTestCase, unittest.TestCase):
+class TestBlockInterleavedLinearOperatorBatch(_BlockSliceTestMixin, LinearOperatorTestCase, unittest.TestCase):
     seed = 0
     should_test_sample = True
 
@@ -49,7 +82,7 @@ def evaluate_linear_op(self, linear_op):
         return actual
 
 
-class TestBlockInterleavedLinearOperatorMultiBatch(LinearOperatorTestCase, unittest.TestCase):
+class TestBlockInterleavedLinearOperatorMultiBatch(_BlockSliceTestMixin, LinearOperatorTestCase, unittest.TestCase):
     seed = 0
     # Because these LTs are large, we'll skil the big tests
     should_test_sample = False