Fix matmul type_inference for 1-D inputs#2719
Open
adityasingh2400 wants to merge 1 commit into
Open
Conversation
matmul.type_inference did not match numpy.matmul (and the op's own value_inference) when either input is 1-D. For a 1-D x it promoted x to [1, K] and then sliced ret_shape[1:] to drop the inserted dimension. The inserted 1 sits at index -2 of the result, not index 0, so this dropped the leading batch dimension instead. For x=(4,), y=(10, 4, 3) it returned (1, 3) where numpy gives (10, 3). A 1-D y was never handled at all: the contraction check reads y_shape[-2] before any promotion, so any 1-D y (e.g. x=(10, 3, 4), y=(4,)) raised IndexError, and the both-1-D case meant to return a scalar crashed the same way. Follow numpy semantics directly: prepend a 1 to a 1-D x, append a 1 to a 1-D y, run the existing batch broadcast, then remove exactly the inserted axes (x's at index -2, y's at index -1, dropping the higher index first). transpose_y on a 1-D y is now rejected to match the existing transpose_x guard. 2-D and batched cases are unchanged. Add a parametrized regression test covering the 1-D x, 1-D y, both 1-D, and transpose combinations. Fixes apple#2263
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.
Summary
matmul.type_inferenceiniOS15/linear.pydoes not matchnumpy.matmul(and the op's ownvalue_inference) when either input is 1-D. This is issue #2263.For the reported case,
x.shape == (4,),y.shape == (10, 4, 3):Root cause
Two separate problems:
1-D
x: the code promotesxto[1, K], computes the broadcast result, then returnsret_shape[1:]to drop the inserted dimension. The inserted1lands at index-2of the result (the result is[*batch, 1, N]), not index0, so the slice drops the leading batch dimension instead.(4,) x (10, 4, 3)returned(1, 3)rather than(10, 3).1-D
y: never handled. The contraction check readsy_shape[-2]before any promotion, so any 1-Dy(e.g.(10, 3, 4) x (4,)) raisedIndexError, and the both-1-D scalar case crashed the same way. This is the second symptom reported in the issue thread.Fix
Follow numpy semantics directly: prepend a
1to a 1-Dx, append a1to a 1-Dy, run the existing batch broadcast, then remove exactly the inserted axes (x's prepended dim at index-2, y's appended dim at index-1, dropping the higher index first so the lower one stays valid).transpose_yon a 1-Dyis now rejected to mirror the existingtranspose_xguard. 2-D and batched cases are untouched.The inferred shape now matches
numpy.matmulfor every combination of{x 1-D, y 1-D, both, neither}x{transpose_x, transpose_y}, including the scalar both-1-D case and symbolic batch dims.Test
Added
TestMatMul::test_builder_1d_input, a parametrized regression test iniOS14/test_linear.pycovering 1-Dx, 1-Dy, both 1-D, and the transpose variants. It fails on the current code (wrong shapes andIndexErrors) and passes with the fix. Existing matmul tests are unaffected.Fixes #2263