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 69ccf8b17a [Relax][Torch] AssertionError: Unsupported function types
['mean.default'] (#18574)
69ccf8b17a is described below
commit 69ccf8b17af764a7c1d636c104d80072a3e42171
Author: Nguyen Duy Loc <[email protected]>
AuthorDate: Mon Dec 22 23:25:15 2025 +0700
[Relax][Torch] AssertionError: Unsupported function types ['mean.default']
(#18574)
## Summary
Happen error when create module from exported_program have torch.mean
without dim.
## Reproduce
- Module:
```
class MeanModule(nn.Module):
def forward(self, x):
return torch.mean(x)
...
# Export → Relax
ep = torch_export(m, (x,))
mod = from_exported_program(ep)
```
- Error log:
```
---------------------------------------------------------------------------
AssertionError Traceback (most recent call last)
Cell In[2], line 13
11 # Export → Relax
12 ep = torch_export(m, (x,))
---> 13 mod = from_exported_program(ep)
15 mod.show()
17 target = "llvm"
File
~/Programming/tvm/python/tvm/relax/frontend/torch/exported_program_translator.py:1783,
in from_exported_program(exported_program, keep_params_as_input,
unwrap_unit_return_tuple, no_bind_return_tuple, run_ep_decomposition)
1780 if run_ep_decomposition:
1781 exported_program = exported_program.run_decompositions()
-> 1783 return ExportedProgramImporter().from_exported_program(
1784 exported_program,
1785 keep_params_as_input,
1786 unwrap_unit_return_tuple,
1787 no_bind_return_tuple,
1788 )
File
~/Programming/tvm/python/tvm/relax/frontend/torch/exported_program_translator.py:1642,
in ExportedProgramImporter.from_exported_program(self, exported_program,
keep_params_as_input, unwrap_unit_return_tuple, no_bind_return_tuple)
1639 nodes: List[fx.Node] = exported_program.graph.nodes
1641 # Find all the missing function types
-> 1642 self._check_unsupported_func_type(nodes)
1644 with self.block_builder.function(
1645 name=func_name, params=list(inputs_vars.values()).copy(),
attrs=func_attrs
1646 ):
1647 output = None
File
~/Programming/tvm/python/tvm/relax/frontend/torch/base_fx_graph_translator.py:182,
in BaseFXGraphImporter._check_unsupported_func_type(self, nodes)
174 def _check_unsupported_func_type(self, nodes: List[fx.Node]):
175 missing_func_types = list(
176 {
177 node.target.__name__
(...) 180 }
181 )
--> 182 assert not missing_func_types, f"Unsupported function types
{missing_func_types}"
AssertionError: Unsupported function types ['mean.default']
```
## Resolve:
- Add "mean.default" into create_convert_map in class
ExportedProgramImporter.
---
.../relax/frontend/torch/exported_program_translator.py | 1 +
.../python/relax/test_frontend_from_exported_program.py | 17 +++++++++++++++++
2 files changed, 18 insertions(+)
diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py
b/python/tvm/relax/frontend/torch/exported_program_translator.py
index 3d6a632fb2..94df0282c8 100644
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -1371,6 +1371,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
"any.dim": self._any,
"any.dims": self._any,
"mean.dim": self._mean,
+ "mean.default": self._mean,
"prod.default": self._prod,
"std.correction": self._std,
"sum.default": self._sum,
diff --git a/tests/python/relax/test_frontend_from_exported_program.py
b/tests/python/relax/test_frontend_from_exported_program.py
index 4a84b50cc9..7894a9fb6d 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -4911,6 +4911,10 @@ def test_mean():
def forward(self, input: torch.Tensor):
return input.mean(-1, keepdim=True)
+ class MeanWithoutDim(Module):
+ def forward(self, input: torch.Tensor):
+ return input.mean()
+
@I.ir_module
class Expected1:
@R.function
@@ -4935,9 +4939,22 @@ def test_mean():
R.output(gv)
return gv
+ @I.ir_module
+ class Expected3:
+ @R.function
+ def main(
+ inp_0: R.Tensor((256, 256), dtype="float32")
+ ) -> R.Tuple(R.Tensor((), dtype="float32")):
+ with R.dataflow():
+ lv: R.Tensor((), dtype="float32") = R.mean(inp_0, axis=None,
keepdims=False)
+ gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv,)
+ R.output(gv)
+ return gv
+
example_args = (torch.randn(256, 256, dtype=torch.float32),)
verify_model(Mean(), example_args, {}, Expected1)
verify_model(MeanKeepDim(), example_args, {}, Expected2)
+ verify_model(MeanWithoutDim(), example_args, {}, Expected3)
def test_sum():
