This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 2fb591c5ba [Relax][PyTorch] Bind symbolic scalar inputs in from_fx
(#19964)
2fb591c5ba is described below
commit 2fb591c5ba4d64f145ca90e946ea374a78fbba8c
Author: Guan-Ming Chiu <[email protected]>
AuthorDate: Thu Jul 9 09:03:35 2026 +0800
[Relax][PyTorch] Bind symbolic scalar inputs in from_fx (#19964)
## Why
- `torch.compile(backend=relax_dynamo(), dynamic=True)` lifts SymInt
scalars as scalar graph inputs
- `from_fx` skips these placeholders, so ops referencing one, e.g.
`view(x.size(0), -1)`, fail with `KeyError`
## How
- Bind sym placeholders to the same-named `tir.Var` from the input
tensors' symbolic shapes; skip as before when none exists
- Add `test_relax_dynamo_dynamic_sym_input_reference`; fails with
`KeyError` without the fix
---
python/tvm/relax/frontend/torch/fx_translator.py | 10 ++++++++--
tests/python/relax/test_frontend_dynamo.py | 20 ++++++++++++++++++++
2 files changed, 28 insertions(+), 2 deletions(-)
diff --git a/python/tvm/relax/frontend/torch/fx_translator.py
b/python/tvm/relax/frontend/torch/fx_translator.py
index 6f40f51f28..75098ee9af 100644
--- a/python/tvm/relax/frontend/torch/fx_translator.py
+++ b/python/tvm/relax/frontend/torch/fx_translator.py
@@ -1093,6 +1093,10 @@ class TorchFXImporter(BaseFXGraphImporter):
# Find all the missing function types
self._check_unsupported_func_type(graph.nodes)
+ from tvm import tirx
+
+ sym_vars = {v.name: v for shape, _ in input_info for v in shape if
isinstance(v, tirx.Var)}
+
with self.block_builder.function(name=func_name, params=inputs.copy(),
attrs=func_attrs):
output = None
with self.block_builder.dataflow():
@@ -1108,11 +1112,13 @@ class TorchFXImporter(BaseFXGraphImporter):
# Translate the model.
for node in graph.nodes:
if node.op == "placeholder":
- assert len(inputs) > 0, "Provided inputs is less than
actual inputs"
if "grapharg" in node.meta and
node.meta["grapharg"].fake_tensor is None:
- # Ignore sym input
+ # Sym input: bind to the matching shape var if
referenced
+ if node.name in sym_vars:
+ self.env[node] = sym_vars[node.name]
continue
+ assert len(inputs) > 0, "Provided inputs is less than
actual inputs"
self.env[node] = inputs.pop(0)
elif node.op == "output":
args = self.retrieve_args(node)
diff --git a/tests/python/relax/test_frontend_dynamo.py
b/tests/python/relax/test_frontend_dynamo.py
index a4a08953b5..ae69905d30 100644
--- a/tests/python/relax/test_frontend_dynamo.py
+++ b/tests/python/relax/test_frontend_dynamo.py
@@ -180,6 +180,26 @@ def test_relax_dynamo_dynamic():
tvm.testing.assert_allclose(opt_func(x, y), opt_func(x, y))
+def test_relax_dynamo_dynamic_sym_input_reference():
+ class ViewModel(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.conv = torch.nn.Conv2d(3, 4, kernel_size=3, padding=1)
+
+ def forward(self, x):
+ return self.conv(x).view(x.size(0), -1)
+
+ model = ViewModel()
+ opt_model = torch.compile(model, backend=relax_dynamo(), dynamic=True)
+
+ with torch.no_grad():
+ for s in (1, 2, 4):
+ inp = torch.randn(s, 3, 8, 8)
+ tvm.testing.assert_allclose(
+ opt_model(inp).detach().numpy(), model(inp).detach().numpy(),
rtol=1e-5, atol=1e-5
+ )
+
+
def test_subgraph_capture():
import torch
