diff --git a/python/tvm/relax/frontend/torch/fx_translator.py b/python/tvm/relax/frontend/torch/fx_translator.py index 6f40f51f2802..75098ee9affc 100644 --- a/python/tvm/relax/frontend/torch/fx_translator.py +++ b/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) diff --git a/tests/python/relax/test_frontend_dynamo.py b/tests/python/relax/test_frontend_dynamo.py index a4a08953b579..ae69905d305e 100644 --- a/tests/python/relax/test_frontend_dynamo.py +++ b/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