[M-Colley]

Fixes cornellius-gp/gpytorch#2747. Alternative to #130 (see below).

Problem

BlockLinearOperator._getitem has a fast path for slices: if all slice bounds are multiples of num_blocks, it divides them by num_blocks and takes 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 inherits it too:

import torch
from linear_operator.operators import BlockDiagLinearOperator

torch.manual_seed(0)
blocks = torch.randn(2, 4, 4)
op = BlockDiagLinearOperator(blocks)  # 8 x 8, two 4 x 4 blocks on the diagonal
print((op[0:4, 0:4].to_dense() - blocks[0]).abs().max())  # tensor(2.1152), should be 0
print(op[0:4, 4:8].to_dense().abs().max())  # tensor(2.1152), off-diagonal block, should be 0

op[0:4, 0:4] returns the top-left 2 x 2 entries of both blocks instead of the first block, with the correct shape, so the result is silently wrong. If the divided slices are not square, it raises base_linear_op must be a batch of square matrices instead (cornellius-gp/gpytorch#2747, e.g. in ExactGP prediction). Before #115 (released in 0.6.1), these slices raised errors; since then the square ones return wrong values. In BoTorch, this made the joint posterior covariance of a ModelListGP with batched multi-output sub-models wrong whenever q was divisible by the number of outputs (meta-pytorch/botorch#3409).

Change

Move the fast path to BlockInterleavedLinearOperator. BlockLinearOperator._getitem keeps the batch-only fast path, so BlockDiagLinearOperator now uses the generic path for slices of the matrix dimensions (as it already did for slices whose bounds are not multiples of num_blocks).

#130 changes the shared base-class fast path to block-diagonal semantics instead. That fixes BlockDiagLinearOperator but breaks BlockInterleavedLinearOperator: for two 4 x 4 interleaved blocks, op[0:4, 0:4] then differs from the dense slice by 2.04. The existing tests do not catch this, because none of the generic test_getitem slices have bounds that are multiples of num_blocks or the block size.

A structure-preserving fast path for slices that select whole diagonal blocks is possible, but base_linear_op[..., i:j, :, :] currently returns the wrong blocks for a CatLinearOperator base when j is the number of blocks (#137), which is how GPyTorch's MultitaskMultivariateNormal.from_independent_mvns builds its covariance. So I left it out of this fix.

Tests

test_getitem_slices compares 100 row x column slice pairs (block-aligned, num_blocks-aligned, straddling, off-diagonal, negative, with a step), with and without batch indices, against the dense matrix for all block-diagonal and interleaved test classes, plus a block-diagonal operator with a CatLinearOperator base.

• main: 16 of the 100 slice pairs fail for each of the 5 block-diagonal test classes (wrong values or errors); the interleaved ones pass.
• Fix BlockLinearOperator._getitem slicing for cross-block indices #130: the 3 interleaved test classes fail (16 slice pairs in total), and so does the CatLinearOperator base (4 slice pairs).
• This PR: all pass. Full suite: 5054 tests pass, except test_psd_safe_cholesky_psd, which also fails on main locally (torch 2.13 on Windows, unrelated).

flake8 passes. Tested with Python 3.13, PyTorch 2.13 (CPU).

🤖 Generated with Claude Code