Skip to content

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

Open
Aharrypotter wants to merge 5 commits intoapache:mainfrom
Aharrypotter:relax-onnx-dynamic-unsqueeze-squeeze-slice
Open

[Relax][ONNX] Fix shape/dynamic restrictions for Squeeze/Unsqueeze and Slice#18955
Aharrypotter wants to merge 5 commits intoapache:mainfrom
Aharrypotter:relax-onnx-dynamic-unsqueeze-squeeze-slice

Conversation

@Aharrypotter
Copy link
Copy Markdown
Contributor

@Aharrypotter Aharrypotter commented Mar 30, 2026

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

@Aharrypotter
[Relax][ONNX] Improve Squeeze/Unsqueeze/Slice handling with dynamic a…
0fec6ab 
…xes support
@Aharrypotter
[Relax][ONNX] Tighten dynamic Unsqueeze/Slice validation and add stru…
5168a67 
…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
@Aharrypotter
[Relax][ONNX] Refactor constant Unsqueeze scalar path and add regress…
8242e43 
…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[bot]
gemini-code-assist bot reviewed Mar 30, 2026
View reviewed changes
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
tlopex requested changes Mar 30, 2026
View reviewed changes
Copy link
Copy Markdown
Member

@tlopex tlopex left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  • 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.

python/tvm/relax/frontend/onnx/onnx_frontend.py Outdated
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}.")
Copy link
Copy Markdown
Member

@tlopex tlopex Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

@Aharrypotter Aharrypotter Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense, I'll fix this.

Copy link
Copy Markdown
Contributor Author

@Aharrypotter Aharrypotter Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@Aharrypotter Aharrypotter force-pushed the relax-onnx-dynamic-unsqueeze-squeeze-slice branch from 4c5965b to 29dc965 Compare March 30, 2026 16:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@tlopex tlopex tlopex requested changes

+1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Reviewers whose approvals may not affect merge requirements

Requested changes must be addressed to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@Aharrypotter @tlopex