Block operator batch scaling broadcasts constants over the internal block axis #135

Bug

Batch-specific scalar multiplication of block operators broadcasts the scalar batch against the internal block axis. This can silently produce wrong matrices, not just shape errors. Reproduced on current main 1c5e24b, Python 3.12.14 / PyTorch 2.10.0+cpu / Windows.

import torch
from linear_operator.operators import DenseLinearOperator, SumBatchLinearOperator

blocks = torch.arange(1., 37.).reshape(3, 3, 2, 2)
op = SumBatchLinearOperator(DenseLinearOperator(blocks))
scale = torch.arange(1., 4.).reshape(3, 1, 1)
expected = op.to_dense() * scale
actual = (op * scale).to_dense()
print((actual - expected).abs().max())  # tensor(88.), expected 0

The same example with BlockDiagLinearOperator or BlockInterleavedLinearOperator has maximum error56. With two batches/three blocks it raises a size mismatch. Two batches/one block also miscomputes SumBatch and raises in the other two operators.

Expected

A scalar for each visible batch should multiply every internal block in that batch by that scalar, matching dense Tensor broadcasting.

Cause / proposed contribution

BlockLinearOperator._mul_constant passes the visible batch scalar directly to a ConstantMulLinearOperator over the base, which has an extra block batch dimension. I am preparing a singleton dimension for that internal block axis, plus dense forward/gradient regressions for all three subclasses, including unequal/equal batch and block counts and broadcast batch scalars.

I checked existing issues/PRs: #25 was about indexing an explicit ConstantMul wrapper; #100 concerns batch summation/permutation with a lazy GPyTorch kernel; #130 changes block slicing. This report exercises direct public multiplication of an already constructed block operator and is separate from those changes.

AI disclosure

OpenAI Codex is autonomously reproducing, implementing and testing this contribution at the account owner's request. No independent human code review is claimed.