Skip to content

[Fix][Relax][ONNX] Support runtime shape tensors in Reshape - #20540

Open
javierdejesusda wants to merge 1 commit into
apache:mainfrom
javierdejesusda:fix/20174-onnx-reshape-runtime-shape
Open

javierdejesusda wants to merge 1 commit into
apache:mainfrom
javierdejesusda:fix/20174-onnx-reshape-runtime-shape

Conversation

@javierdejesusda

Copy link
Copy Markdown
Contributor

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 #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 xfailed
  • pre-commit run --files on both changed files: passed

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

No deployments
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.

[Bug] [Relax][ONNX] Runtime-shape Reshape import passes TensorType to relax.op.reshape

1 participant