Skip to content

Fix shared vector backward for batched operator matmul - #154

Open
M-Colley wants to merge 1 commit into
cornellius-gp:mainfrom
M-Colley:codex/fix-batched-vector-matmul-backward
Open

M-Colley wants to merge 1 commit into
cornellius-gp:mainfrom
M-Colley:codex/fix-batched-vector-matmul-backward

Conversation

@M-Colley

Copy link
Copy Markdown

Multiplying a batched operator by a shared one-dimensional vector produces a batched vector output. Matmul.backward currently uses the dimensionality of the output gradient to decide whether it is a vector, so the operator batch axes cause it to take the matrix branch. Differentiating the matrix, the vector, or both can then fail with incompatible matrix dimensions.

For example, (DenseLinearOperator(torch.randn(2, 3, 4, requires_grad=True)) @ torch.randn(4, requires_grad=True)).sum().backward() raises even though the forward result agrees with torch.matmul.

Use the original saved RHS shape to establish vector semantics, normalize the RHS and incoming gradient to matrix operands for both VJPs, then remove the final singleton column before the existing batch reduction. The original unbatched CUDA vector path still uses a two-dimensional operand to avoid the existing GEMV issue. Forward values and the matrix-RHS path are unchanged.

Validation against main 1c5e24b:

  • New regressions fail in 34 CPU/CUDA subcases before the repair; 12 unbatched controls already pass.
  • Torch 2.13 CPU/CUDA: 483 tests and 46 subcases pass across the new regressions and existing matmul/Cat suites.
  • Torch 2.14.1 CPU: 483 tests and 23 subcases pass.
  • Checks include rectangular/multiple operator batch axes, singleton row/column dimensions, every matrix/vector trainability combination, noncontiguous shared vectors, cached/reconstructed operator modes, structured batched operators, mixed first/second/third derivatives, gradcheck and gradgradcheck.
  • ufmt check, flake8, and git diff --check pass for the changed files.

This repair is independent of #153 (concatenated operator parameter higher derivatives): the structured-operator regression here freezes the operator and tests the shared vector's graph, so no Cat derivative change is needed.

AI assistance: implemented and tested with OpenAI Codex. The reproductions, regression runs, and formatting checks above were executed locally.

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