[DiadkoShmek]

Block-diagonal/diagonal multiplication reshapes the diagonal operand using the block operand's batch axes. Broadcast-compatible prefixes can therefore fail to reshape or silently move scenarios to the wrong axis. For example:

import torch
from linear_operator.operators import BlockDiagLinearOperator, DenseLinearOperator, DiagLinearOperator

blocks = torch.eye(2).repeat(2, 1, 3, 1, 1)
block = BlockDiagLinearOperator(DenseLinearOperator(blocks))
diag = DiagLinearOperator(torch.arange(1., 13.).reshape(1, 2, 6))
print((block @ diag).shape)  # currently (2, 1, 6, 6); expected (2, 2, 6, 6)
print((diag @ block).shape)  # same batch-axis problem

Split only the diagonal event axis into (num_blocks, block_size), retaining that operand's own batch prefix. The existing block multiplication then performs normal broadcasting. This requires one reshape correction in each multiplication order, without materializing the full matrix. A collapsed covariance batch can also produce an incorrect Gaussian density when a public distribution broadcasts it against its mean.

Validation is CPU-only:

• The complete BlockDiag and Diag test modules pass 510 methods on PyTorch 2.4.1/NumPy 2.0.2, and separately on the configured stable PyTorch 2.0.0/NumPy 1.26.4 environment. The new regression fails on BASE in eight subtests and raises in eight others; controls remain passing.
• A separate NumPy/SPD/GP-score oracle passes 80 configurations/640 comparisons. Independent review passes 112 product configurations and 64 well-conditioned Gaussian/default-interleaved multitask configurations (880 value/first-gradient comparisons), including equal-cardinality crossed batch axes. Four incompatible-batch/event rejection controls are unchanged.
• Exact workflow Flake8 5.0.4 and pre-commit pass; full all-files hooks and local wheel/sdist builds also pass on BASE and candidate. Both HTML documentation builds succeed, with zero Sphinx logger warnings but the same 25 existing Python RuntimeWarnings from the unchanged type formatter in docs/source/conf.py. No warning suppression or warning-free claim.

This preserves supported operand broadcasting; it does not redesign incompatible dimensions or arbitrary reshapes. No full repository/nightly numerical suite, GPU, complex/empty-event support, higher AD or performance improvement is claimed. Known noninterleaved multitask-density mapping behavior is outside this repair; the public Gaussian checks use the default interleaved path. The separate triangular-solve change in #158 is not included.

I directed this work and take responsibility for the contribution and its follow-up. AI assistance was used in producing the change.