Repository navigation
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Multiplying a batched operator by a shared one-dimensional vector produces a batched vector output.
Matmul.backwardcurrently 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 withtorch.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:ufmt check,flake8, andgit diff --checkpass 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.