[Fix][Relax][ONNX] Support runtime shape tensors in Reshape - #20540
Open
javierdejesusda wants to merge 1 commit into
Open
javierdejesusda wants to merge 1 commit into
javierdejesusda wants to merge 1 commit into
Conversation
The Relax ONNX importer handled `Reshape` shape inputs that are constants or `Shape` outputs. When the shape is a runtime int64 tensor, such as a graph input, the converter passed the tensor to `relax.op.reshape`, which fails with `TypeError: Reshape requires the input new shape to be Shape. However, the given one is relax.TensorType`. For a runtime shape, the converter now builds the target shape with Relax tensor ops and converts it to a `ShapeExpr` with `_tensor_to_shape_expr`, the way the dynamic-axes path of `Unsqueeze` does. With `allowzero=0` a 0 entry copies the input dimension at the same index, and a -1 entry is the number of input elements divided by the product of the other entries, so neither reaches `relax.op.reshape` as a literal. The copy index is clamped to the last input dimension so `take` stays in bounds, and a scalar input skips the copy. With `allowzero=1` a literal 0 stays 0. The constant and `ShapeExpr` paths are unchanged. The new tests make the shape a graph input and compare the imported model with onnxruntime for the shape from the issue, `-1`, `0`, a rank increase and decrease, a scalar input, a symbolic batch dimension, and `allowzero=1` including a zero-size input. One more test checks the error for an unknown shape length, and one checks that a `Shape` output of an unknown-rank input still imports. The 16 runtime-shape tests fail on `main` with the `TypeError` above. Fixes apache#20174 Related to apache#17892, which reports the same error on v0.20.0. Tests: - `pytest -q -n 6 tests/python/relax/test_frontend_onnx.py`: 633 passed, 7 skipped, 4 xfailed - `pre-commit run --files` on both changed files: passed
This branch has not been deployed
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.
The Relax ONNX importer handled
Reshapeshape inputs that are constants orShapeoutputs. When the shape is a runtime int64 tensor, such as a graph input, the converter passed the tensor torelax.op.reshape, which fails withTypeError: Reshape requires the input new shape to be Shape. However, the given one is relax.TensorType.For a runtime shape, the converter now builds the target shape with Relax tensor ops and converts it to a
ShapeExprwith_tensor_to_shape_expr, the way the dynamic-axes path ofUnsqueezedoes. Withallowzero=0a 0 entry copies the input dimension at the same index, and a -1 entry is the number of input elements divided by the product of the other entries, so neither reachesrelax.op.reshapeas a literal. The copy index is clamped to the last input dimension sotakestays in bounds, and a scalar input skips the copy. Withallowzero=1a literal 0 stays 0. The constant andShapeExprpaths are unchanged.The new tests make the shape a graph input and compare the imported model with onnxruntime for the shape from the issue,
-1,0, a rank increase and decrease, a scalar input, a symbolic batch dimension, andallowzero=1including a zero-size input. One more test checks the error for an unknown shape length, and one checks that aShapeoutput of an unknown-rank input still imports. The 16 runtime-shape tests fail onmainwith theTypeErrorabove.Fixes #20174
Related to #17892, which reports the same error on v0.20.0.
Tests:
pytest -q -n 6 tests/python/relax/test_frontend_onnx.py: 633 passed, 7 skipped, 4 xfailedpre-commit run --fileson both changed files: passed