From ed4238f64b3080a6e7720fdf94d0b3043527f624 Mon Sep 17 00:00:00 2001 From: Mark Colley <39366189+M-Colley@users.noreply.github.com> Date: Mon, 12 Oct 2026 08:47:54 +0100 Subject: [PATCH] Fix slicing of BlockDiagLinearOperator BlockLinearOperator._getitem had a fast path for slices: if all slice bounds are multiples of num_blocks, it divided them by num_blocks and took those rows and columns from every block. That is right for BlockInterleavedLinearOperator, where row i of the matrix is row i // num_blocks of block i % num_blocks, but BlockDiagLinearOperator inherited it too. For example, with two 4 x 4 blocks, op[..., 0:4, 0:4] returned the top-left 2 x 2 entries of both blocks instead of the first block, and op[..., 0:4, 4:8] returned non-zero entries. If the divided slices were not square, it raised "base_linear_op must be a batch of square matrices" instead (cornellius-gp/gpytorch#2747). Move the fast path to BlockInterleavedLinearOperator. BlockLinearOperator keeps the batch-only fast path, so BlockDiagLinearOperator now uses the generic path for slices of the matrix dimensions. Add tests that compare slices of block-diagonal and interleaved operators against their dense matrices, including slices whose bounds are multiples of num_blocks or of the block size. Co-Authored-By: Claude Opus 5.5 --- .../block_interleaved_linear_operator.py | 24 +++++++++ .../operators/block_linear_operator.py | 33 ++---------- .../test_block_diag_linear_operator.py | 54 +++++++++++++++++-- .../test_block_interleaved_linear_operator.py | 39 ++++++++++++-- 4 files changed, 113 insertions(+), 37 deletions(-) diff --git a/linear_operator/operators/block_interleaved_linear_operator.py b/linear_operator/operators/block_interleaved_linear_operator.py index 69f0429a..e7ad38a8 100644 --- a/linear_operator/operators/block_interleaved_linear_operator.py +++ b/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) diff --git a/linear_operator/operators/block_linear_operator.py b/linear_operator/operators/block_linear_operator.py index 4a8e9e51..f9af4bad 100644 --- a/linear_operator/operators/block_linear_operator.py +++ b/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) diff --git a/test/operators/test_block_diag_linear_operator.py b/test/operators/test_block_diag_linear_operator.py index fb6b8f1a..90682c8f 100644 --- a/test/operators/test_block_diag_linear_operator.py +++ b/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 diff --git a/test/operators/test_block_interleaved_linear_operator.py b/test/operators/test_block_interleaved_linear_operator.py index 27baacbb..f496fb19 100644 --- a/test/operators/test_block_interleaved_linear_operator.py +++ b/test/operators/test_block_interleaved_linear_operator.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 import unittest +from itertools import product import torch @@ -8,7 +9,39 @@ 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