[Relax][ONNX] Fix shape/dynamic restrictions for `Squeeze`/`Unsqueeze` and `Slice` by Aharrypotter · Pull Request #18955 · apache/tvm

Aharrypotter · 2026-03-30T09:41:40Z

Summary

Relates to #18945.

This PR improves ONNX frontend handling for dynamic Unsqueeze/Squeeze/Slice, tightens validation paths, and adds targeted structural/negative regression tests.

Refactor constant-path Unsqueeze lowering to use a single reshape based on computed target shape.
Remove scalar-specific branching and repeated expand_dims in the constant path.
Add/keep structural helper usage in ONNX frontend tests for Relax call-op checks.
Add regression coverage for scalar-input Unsqueeze.

Changes

Add dynamic-axes conversion paths for Unsqueeze and Squeeze:
- infer output shape via runtime shape-tensor construction
- lower to relax.reshape with validated shape rank/length assumptions
Improve Slice conversion robustness:
- support dynamic parameter forms with stricter rank/length validation
- reject invalid zero-step inputs when statically known
- fix docstring wording (Splice -> Slice)
Strengthen ONNX frontend tests:
- negative test for duplicate Unsqueeze axes
- structural IR check for dynamic Slice (relax.dynamic_strided_slice present, relax.strided_slice absent)
- negative test for zero-step Slice
Refactor constant-path Unsqueeze scalar handling:
- replace scalar special-casing + repeated expand_dims with one target-shape reshape
- add scalar-input regression test
Restore shared test helper used by structural Relax call-op checks.

validation

ruff check: passed
pre-commit --files: passed
pytest: 8 passed

…xes support

…ctural tests - fix Slice converter docstring typo ("Splice" -> "Slice") for consistency - add explicit validation to reject zero step values in Slice for both constant and dynamic-parameter paths - add Unsqueeze negative test to reject duplicate axes - strengthen structural IR test for dynamic Slice to assert relax.dynamic_strided_slice is used and relax.strided_slice is not - add Slice negative test for zero-step input Validation: - python -m ruff check python/tvm/relax/frontend/onnx/onnx_frontend.py tests/python/relax/test_frontend_onnx.py - python -m pre_commit run --files python/tvm/relax/frontend/onnx/onnx_frontend.py tests/python/relax/test_frontend_onnx.py - python -m pytest -n 1 tests/python/relax/test_frontend_onnx.py -k "unsqueeze_dynamic_axes or unsqueeze_duplicate_axes_validation or slice_dynamic_inputs_ir or slice_dynamic_inputs_length_validation or slice_zero_step_validation" -v Result: - 7 passed

…ion test - refactor constant Unsqueeze lowering to build target shape then reshape instead of scalar special-casing plus repeated expand_dims - add Unsqueeze scalar-input regression test to cover the new path - restore helper used by structural ONNX frontend IR checks Validation: - python -m ruff check python/tvm/relax/frontend/onnx/onnx_frontend.py tests/python/relax/test_frontend_onnx.py - python -m pre_commit run --files python/tvm/relax/frontend/onnx/onnx_frontend.py tests/python/relax/test_frontend_onnx.py - python -m pytest -n 1 tests/python/relax/test_frontend_onnx.py -k "unsqueeze_scalar_input or unsqueeze_dynamic_axes or unsqueeze_duplicate_axes_validation or slice_dynamic_inputs_ir or slice_dynamic_inputs_length_validation or slice_zero_step_validation" -v Result: - 8 passed

gemini-code-assist

Code Review

This pull request enhances the ONNX frontend for TVM Relax by implementing support for dynamic axes and parameters in the Unsqueeze, Squeeze, and Slice operators. It introduces several utility functions for handling runtime shape calculations and tensor rank validation, such as _get_known_tensor_rank and _build_unsqueezed_shape_tensor. The changes also include comprehensive new test cases to verify dynamic behavior and input validation. Feedback was provided to simplify the dimension-checking logic in the _get_known_tensor_length helper function to improve readability.

tlopex

Please fix the Squeeze behavior for ShapeExpr input when axis=None. Per ONNX spec, it should remove all size-1 dimensions, but the current path returns the input unchanged when len(data) > 1.
Please take a look at the Slice dynamic path. It currently converts ShapeExpr to a tensor and returns a tensor from dynamic_strided_slice, while the constant path returns ShapeExpr or relax.const. That type mismatch may cause downstream issues.
_get_known_tensor_length could also be simplified a bit — the second if struct_info.ndim != 1 is effectively just the ndim == -1 case, so using if struct_info.ndim == -1: return None would be clearer.

tlopex · 2026-03-30T11:56:48Z

+        if axis < 0:
+            axis += rank
+        if axis < 0 or axis >= rank:
+            raise ValueError(f"{op_name} axis {axis} is out of range for rank {rank}.")


Please preserve the original axis value in _normalize_constant_axes for the error message. Right now it reports the normalized value instead of the user input, which can be confusing.

Makes sense, I'll fix this.

Hi @tlopex , thank you for the detailed review! I've updated the PR to address your and gemini-code-assist's feedback

Especially, regarding the Slice dynamic path: I've added an explicit check to reject ShapeExpr inputs. This prevents potential type mismatches between constant and dynamic paths while keeping the implementation clean.

…lidation

The previous conflict merge incorrectly changed the return type of collect_relax_call_ops from list to set. This caused test failures because: 1. Multiple tests rely on .count() to verify the exact number of operator insertions (e.g., ensuring both inputs of MatMulInteger16 are casted). 2. Several tests use '== []' to assert that no operators were generated (verifying constant folding), which is type-incompatible with set(). Restoring list semantics ensures both structural validation and type consistency in tests.

Aharrypotter · 2026-03-31T05:58:57Z

Summary of the fix for `collect_relax_call_ops` merge conflict.

During the recent conflict merge, the helper function collect_relax_call_ops in test_frontend_onnx.py was inadvertently changed to return a set instead of a list. This broke several structural IR tests that require:

Count Verification: Some tests (like test_matmulinteger16_ir) need to verify that a specific operator appears exactly $N$ times (e.g., call_ops.count("relax.astype") == 2). A set deduplicates these entries, losing this critical information.
Type Consistency: Tests that assert collect_relax_call_ops(...) == [] fail when the function returns an empty set(), as the types do not match.
Order Preservation: Maintaining the post-order traversal sequence is better for future-proofing graph topology assertions.

This commit restores the list[str] return type and updates the implementation in test_frontend_onnx.py to ensure all tests pass correctly.

Aharrypotter · 2026-03-31T07:26:22Z

cc @tlopex

tlopex

LGTM

Aharrypotter added 3 commits March 30, 2026 13:38

[Relax][ONNX] Improve Squeeze/Unsqueeze/Slice handling with dynamic a…

0fec6ab

…xes support

gemini-code-assist Bot reviewed Mar 30, 2026

View reviewed changes

Comment thread python/tvm/relax/frontend/onnx/onnx_frontend.py Outdated

tlopex requested changes Mar 30, 2026

View reviewed changes

[Relax][ONNX] tighten Slice dynamic checks and improve axis/length va…

29dc965

…lidation

Aharrypotter force-pushed the relax-onnx-dynamic-unsqueeze-squeeze-slice branch from 4c5965b to 29dc965 Compare March 30, 2026 16:44

Aharrypotter added 2 commits March 31, 2026 00:46

Merge branch 'main' into relax-onnx-dynamic-unsqueeze-squeeze-slice

4eb71c6

tlopex approved these changes Mar 31, 2026

View reviewed changes

tlopex merged commit fd9d9db into apache:main Mar 31, 2026
9 checks passed

Aharrypotter deleted the relax-onnx-dynamic-unsqueeze-squeeze-slice branch April 4, 2026 07:05

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Relax][ONNX] Fix shape/dynamic restrictions for `Squeeze`/`Unsqueeze` and `Slice`#18955

[Relax][ONNX] Fix shape/dynamic restrictions for `Squeeze`/`Unsqueeze` and `Slice`#18955
tlopex merged 6 commits intoapache:mainfrom
Aharrypotter:relax-onnx-dynamic-unsqueeze-squeeze-slice

Aharrypotter commented Mar 30, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

tlopex left a comment

Uh oh!

tlopex Mar 30, 2026

Uh oh!

Aharrypotter Mar 30, 2026

Uh oh!

Aharrypotter Mar 30, 2026

Uh oh!

Aharrypotter commented Mar 31, 2026

Uh oh!

Aharrypotter commented Mar 31, 2026

Uh oh!

tlopex left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

Aharrypotter commented Mar 30, 2026

Summary

Changes

validation

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

tlopex left a comment

Choose a reason for hiding this comment

Uh oh!

tlopex Mar 30, 2026

Choose a reason for hiding this comment

Uh oh!

Aharrypotter Mar 30, 2026

Choose a reason for hiding this comment

Uh oh!

Aharrypotter Mar 30, 2026

Choose a reason for hiding this comment

Uh oh!

Aharrypotter commented Mar 31, 2026

Summary of the fix for collect_relax_call_ops merge conflict.

Uh oh!

Aharrypotter commented Mar 31, 2026

Uh oh!

tlopex left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Summary of the fix for `collect_relax_call_ops` merge conflict.