[robrui]

See cornellius-gp/gpytorch#2747 for the full analysis and reproduction.

The modulo check used num_blocks instead of block_size, so slices aligned to individual blocks were incorrectly handled. Cross-block slices fell through to block-level indexing instead of the generic path, producing non-square sub-blocks.

The fix switches to block_size for alignment checks, detects cross-block slices to fall back to the generic path, and selects blocks from the batch dimension rather than passing block indices as matrix row/col indices.

[M-Colley]

Thanks for picking this up! One concern: BlockLinearOperator._getitem is shared by BlockDiagLinearOperator and BlockInterleavedLinearOperator, and the num_blocks-based fast path is actually correct for the interleaved layout (row i of the matrix is row i // num_blocks of block i % num_blocks). Switching it to block_size fixes the block-diagonal case but breaks the interleaved one:

import torch
from linear_operator.operators import BlockInterleavedLinearOperator, DenseLinearOperator

torch.manual_seed(0)
op = BlockInterleavedLinearOperator(DenseLinearOperator(torch.randn(2, 4, 4)))  # 8 x 8
print((op[0:4, 0:4].to_dense() - op.to_dense()[0:4, 0:4]).abs().max())
# main: tensor(0.)   this PR: tensor(2.0447)

The existing test suite passes with this PR because none of the generic test_getitem slices (0:2, 1:, ...) have bounds that are multiples of num_blocks or the block size, and the new test only checks shapes. The block-diagonal bug also has a silent variant: when the divided slices happen to be square, e.g. op[0:4, 0:4] for two 4 x 4 blocks, main returns the top-left 2 x 2 entries of both blocks instead of the first block, with the correct shape.

Also, with a CatLinearOperator base (as built by GPyTorch's MultitaskMultivariateNormal.from_independent_mvns), selecting blocks via base[..., block_slice] runs into #137 when the slice ends at the last block, e.g. op[4:12, 4:12] with three 4 x 4 blocks raises a shape error.

I opened #142, which moves the fast path to BlockInterleavedLinearOperator and lets BlockDiagLinearOperator use the generic path, with value-based tests for both layouts (they fail on main for the block-diagonal operators and on this PR for the interleaved ones). Happy to close it in favor of this one if you'd rather update this PR instead.