Skip to content

[Fix][Relax] Raise error on non-unit dim ONNX Squeeze axis - #20188

Merged
tlopex merged 2 commits into
apache:mainfrom
OmarAzizi:fix-onnx-non-unit-axis
Aug 27, 2026
Merged

tlopex merged 2 commits into
apache:mainfrom
OmarAzizi:fix-onnx-non-unit-axis

Conversation

@OmarAzizi

@OmarAzizi OmarAzizi commented Aug 25, 2026 •

Copy link
Copy Markdown
Contributor

TVM's ONNX Squeeze converter accepted axes with non-unit dimensions instead of raising an error silently passing tensors through unchanged. This adds a check in the ONNX frontend that raises ValueError when a squeeze axis's statically known size isn't 1, matching the ONNX spec.

Fixes #20185

Edit: The code merged from #20147 was causing a linter error so Claude assisted me in reformatting the code

return
for axis in _normalize_constant_axes(list(axes), rank, "Squeeze"):
extent = ty.shape.values[axis]
if isinstance(extent, tirx.IntImm) and int(extent.value) != 1:

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

This only rejects statically known non-unit extents, so the same ONNX violation remains for symbolic dimensions. For an input declared as (N, 3) with constant axes [0], running with shape (2, 3) correctly raises in ONNX Runtime, while TVM imports it and silently returns a (3,) view containing the first row. Please add a runtime check that selected extents equal 1, or reject the import when this cannot be proven, and add a symbolic-dimension regression test. The runtime-axes path needs the same validation.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

You're right, thank you. I went with the second approach because it's similar to how other parts in the frontend handle these cases.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Thanks for handling symbolic dimensions. Runtime axes still bypass this check though, so invalid non-unit axes remain accepted. Please add a runtime assertion or reject non-constant axes during import, with a regression test.

@OmarAzizi OmarAzizi Aug 27, 2026 •

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

You're welcome.

Correct me if I'm worng, but regarding a runtime assertion I don't see assert_op used in any frontend and match_cast only used to bind symbolic vars, never to assert a value. The usual pattern is to reject an import when a value isn't a compile-time constant. Should I just do that here as well?

@tlopex

tlopex commented Aug 27, 2026

Copy link
Copy Markdown
Member

you can rebase the main since the lint issue is fixed

@OmarAzizi
OmarAzizi force-pushed the fix-onnx-non-unit-axis branch from a920415 to e8afc7e Compare August 27, 2026 20:03
@tlopex

tlopex commented Aug 27, 2026

Copy link
Copy Markdown
Member

LGTM

@tlopex
tlopex merged commit c9e9b85 into apache:main Aug 27, 2026
8 checks passed
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][Frontend][ONNX] Squeeze silently ignores axes that select non-unit dimensions instead of raising an error (violates ONNX spec)

2 participants