Skip to content

Fix rectangular batch-repeat matrix multiplication - #157

Open
DiadkoShmek wants to merge 1 commit into
cornellius-gp:mainfrom
DiadkoShmek:lezo2-fix-rectangular-batch-repeat-matmul
Open

DiadkoShmek wants to merge 1 commit into
cornellius-gp:mainfrom
DiadkoShmek:lezo2-fix-rectangular-batch-repeat-matmul

Conversation

@DiadkoShmek

Copy link
Copy Markdown

Problem

Batch-repeated rectangular operators multiply the unrepeated base directly, so repeating a non-singleton batch can raise a batch-size mismatch even when dense Torch multiplication is valid:

import torch
from linear_operator.operators import DenseLinearOperator

x = torch.arange(24, dtype=torch.double).reshape(2, 3, 4)
rhs = torch.arange(80, dtype=torch.double).reshape(4, 4, 5)
actual = DenseLinearOperator(x).repeat(2, 1, 1) @ rhs
expected = x.repeat(2, 1, 1) @ rhs

Change

Use the existing repeat-batches-to-columns algorithm for rectangular multiplication as well. Fold the broadcast input shape (..., N, C), apply the base operator, then unfold the output shape (..., M, C). Explicit column sizes retain empty products, and vector normalization preserves direct vector calls. This retains structured/factorized multiplication without materializing the repeated matrix.

Validation

  • New failing-first tests: two original failures and one passing unbatched control.
  • Full owner module: 230 passed, including existing square solve controls and new batched rectangular, empty-product, vector and factorized-gradient regressions.
  • Separate root review: 96 factorized public products, all failing on the original owner and passing after the change; 288 gradient tensor comparisons, transpose/noncontiguous/broadcast cases, five empty products, and eight solve/Cholesky helper controls. Dense-materialization guards cover forward and backward.
  • A broader builder oracle checks 352 products, 44 vector-forward cases and 696 gradient tensors.
  • Repository-pinned ufmt/Black/usort and flake8 checks pass; git diff --check passes.

CPU Python 3.13.7 / Torch 2.11 only; no full-repository, documentation or GPU/version-matrix run. Existing generic zero-width/broadcast RHS gradient and batched-vector autograd limitations are unchanged and reproduced separately with plain DenseLinearOperator.

AI assistance: OpenAI Codex investigated, implemented and independently reviewed this fix. No human review or upstream approval is claimed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant