[Relax][PyTorch] Bind symbolic scalar inputs in from_fx

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

@@ -1093,6 +1093,10 @@ def from_fx(
         # 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 @@ def from_fx(
                 # 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)

tests/python/relax/test_frontend_dynamo.py

@@ -180,6 +180,26 @@ def Func1(x, y):
             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