[Relax][PyTorch] Support basic range constraints #18429

guan404ming · 2025-11-10T12:41:31Z

Related Issue

related to #17818

Why

Add support for PyTorch range constraints in Relax frontend
Enables proper handling of dynamic shapes with min/max bounds

How

Extract range constraints from PyTorch ExportedProgram and store in shape_var_constraints function attribute
Add test test_dynamic_shape_with_range_constraints to verify constraint extraction works correctly

guan404ming · 2025-11-10T16:36:37Z

cc @mshr-h

mshr-h

Thanks! Few changes are needed.

mshr-h · 2025-11-11T10:04:44Z

python/tvm/relax/frontend/torch/exported_program_translator.py

-                else s
-                for s in torch_shape
-            ]
+            # Create TIR variables for symbolic dimensions


Looks like there's no functional changes. Any reason for the change?

Sorry about this, it seems like legacy change in #17898. Let me revert this.

mshr-h · 2025-11-11T10:09:59Z

python/tvm/relax/frontend/torch/exported_program_translator.py

+        if range_constraints:
+            if func_attrs is None:
+                func_attrs = {}
+            func_attrs["shape_var_constraints"] = range_constraints


Please use tir_var_upper_bound to annotate upper bound.
I grepped the tvm code base and I found that there's no lower bound annotation. So I don't think we need to keep it at the moment. If we have a real use case for it, it's fine to keep it in the Relax module.
https://github.com/apache/tvm/blob/main/src/relax/transform/static_plan_block_memory.cc#L62-L66

Thanks for letting me know, I've updated to use tir_var_upper_bound to annotate upper bound.

mshr-h · 2025-11-11T10:10:33Z

tests/python/relax/test_frontend_from_exported_program.py

+    mod = from_exported_program(exported_program)
+
+    main_func = mod["main"]
+    assert hasattr(main_func, "attrs"), "Function should have attributes"


Please use structual equality instead of manually checking the attributes.

Sure, I've updated the test to follow the convention.

gemini-code-assist

Code Review

This pull request adds support for PyTorch range constraints in the Relax frontend. The changes correctly extract range constraints from the ExportedProgram and attach them as function attributes. A new test case is included to verify this functionality. My review includes a couple of suggestions to improve code conciseness and test robustness.

python/tvm/relax/frontend/torch/exported_program_translator.py

tests/python/relax/test_frontend_from_exported_program.py

mshr-h · 2025-11-11T10:44:34Z

tests/python/relax/test_frontend_from_exported_program.py

+    dynamic_shapes = {"x1": {0: batch}, "x2": {0: batch}}
+    exported_program = export(DynamicModel(), args=example_args, dynamic_shapes=dynamic_shapes)
+
+    mod = from_exported_program(exported_program)


Should be from_exported_program(exported_program, run_ep_decomposition=True)
See #18399

Thanks for providing useful info!

Should we also need to update tests in https://github.com/apache/tvm/blob/f574031657faa035ffd1875050316b11821187cf/tests/python/relax/test_frontend_from_exported_program.py?

I could help with this if needed.

Of course you can help update it!

Sure! I would open another PR to handle this, thanks for your reply.

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

mshr-h

LTGM. Thanks!

guan404ming · 2025-11-12T08:54:32Z

Thanks for your detailed comments and reviews!

Support basic range constraints

77685df

guan404ming force-pushed the fix-relax-pytorch-constraints branch from 05fd731 to 77685df Compare November 10, 2025 13:52

guan404ming marked this pull request as ready for review November 10, 2025 16:30

mshr-h self-requested a review November 11, 2025 01:07

mshr-h reviewed Nov 11, 2025

View reviewed changes

gemini-code-assist bot reviewed Nov 11, 2025

View reviewed changes

python/tvm/relax/frontend/torch/exported_program_translator.py Outdated Show resolved Hide resolved

tests/python/relax/test_frontend_from_exported_program.py Outdated Show resolved Hide resolved

mshr-h requested changes Nov 11, 2025

View reviewed changes

guan404ming and others added 2 commits November 11, 2025 23:39

Apply gemini-code-assist suggestions

1ad828f

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Apply reviewer comments

34476a4

guan404ming force-pushed the fix-relax-pytorch-constraints branch from 93b699f to 34476a4 Compare November 11, 2025 16:32

guan404ming added 2 commits November 12, 2025 00:51

Fix lint error

56e7792

Refactor frontend test to use consistent size variable

2921bf0

guan404ming requested a review from mshr-h November 12, 2025 04:27

mshr-h approved these changes Nov 12, 2025

View reviewed changes

mshr-h merged commit 394f668 into apache:main Nov 12, 2025
14 checks passed

This was referenced Nov 12, 2025

[Tracking Issue] [Relax][PyTorch] Dynamic shape support in ExportedProgram #17818

Closed

feat(relax/frontend/torch): Add basic range constraint support #17898

Closed

guan404ming deleted the fix-relax-pytorch-constraints branch November 12, 2025 08:54

This was referenced Nov 14, 2025

[Relax][PyTorch] Add lower bound support for range constraints #18447

Merged

[Relax][PyTorch] Support advanced range constraints (addition) #18452

Merged

guan404ming changed the title ~~[Relax][Pytorch] Support basic range constraints~~ [Relax][PyTorch] Support basic range constraints Nov 14, 2025

guan404ming mentioned this pull request Nov 15, 2025

[Relax][PyTorch] Add decomposed operator support for Binary #18458

Merged

[Relax][PyTorch] Support basic range constraints #18429

[Relax][PyTorch] Support basic range constraints #18429

Uh oh!

Conversation

guan404ming commented Nov 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Related Issue

Why

How

Uh oh!

guan404ming commented Nov 10, 2025

Uh oh!

mshr-h left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mshr-h Nov 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

mshr-h Nov 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mshr-h left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

guan404ming commented Nov 12, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

guan404ming commented Nov 10, 2025 •

edited

Loading

mshr-h Nov 11, 2025 •

edited

Loading

mshr-h Nov 11, 2025 •

edited

Loading