gemini-code-assist[bot] commented on code in PR #18544:
URL: https://github.com/apache/tvm/pull/18544#discussion_r2588227508
##########
python/tvm/relax/frontend/torch/exported_program_translator.py:
##########
@@ -1604,9 +1603,15 @@ def from_exported_program(
keep_params_as_input: bool,
unwrap_unit_return_tuple: bool,
no_bind_return_tuple: bool,
+ custom_convert_map: Dict[str, Callable[[BaseFXGraphImporter, fx.Node],
relax.Var]],
Review Comment:
There is a mismatch between the type hint for `custom_convert_map` and how
the callable is used. The type hint `Callable[[BaseFXGraphImporter, fx.Node],
relax.Var]` implies a signature of `(importer, node)`. However, the call at
line 1680 is `self.convert_map[func_name](node, self)`, which implies a
signature of `(node, importer)`. This will cause a runtime error due to
mismatched argument types. Please correct the order of arguments in the type
hint.
```suggestion
custom_convert_map: Dict[str, Callable[[fx.Node,
BaseFXGraphImporter], relax.Var]],
```
##########
python/tvm/relax/frontend/torch/base_fx_graph_translator.py:
##########
@@ -45,6 +45,17 @@ def __init__(self) -> None:
########## Utilities ##########
+ def update_convert_map(self, custom_convert_map: dict):
Review Comment:
The type hint for `custom_convert_map` is `dict`, which is too generic.
Please use a more specific type hint like `Dict[str, Callable]` for better type
safety and clarity. The necessary types are already imported in this file.
```suggestion
def update_convert_map(self, custom_convert_map: Dict[str, Callable]):
```
##########
python/tvm/relax/frontend/torch/base_fx_graph_translator.py:
##########
@@ -171,12 +182,16 @@ def _retrieve_args(self, node):
else:
return node
- def _check_unsupported_func_type(self, nodes: List[fx.Node]):
+ def _check_unsupported_func_type(self, nodes: List[fx.Node],
custom_convert_map: Dict[str, Callable] = None):
+ supported_op_list = set(self.convert_map.keys())
+ if custom_convert_map:
+ supported_op_list =
supported_op_list.union(custom_convert_map.keys())
+
missing_func_types = list(
{
node.target.__name__
for node in nodes
- if node.op == "call_function" and node.target.__name__ not in
self.convert_map
+ if node.op == "call_function" and node.target.__name__ not in
supported_op_list
}
)
assert not missing_func_types, f"Unsupported function types
{missing_func_types}"
Review Comment:
The `custom_convert_map` parameter is redundant here. The `self.convert_map`
is updated with custom ops before this method is called in both
`fx_translator.py` and `exported_program_translator.py`. You can simplify this
method by removing the `custom_convert_map` parameter and directly using
`self.convert_map`.
This change will also require updating the call site in
`python/tvm/relax/frontend/torch/exported_program_translator.py` at line 1646
to `self._check_unsupported_func_type(nodes)`.
```python
def _check_unsupported_func_type(self, nodes: List[fx.Node]):
missing_func_types = list(
{
node.target.__name__
for node in nodes
if node.op == "call_function" and node.target.__name__ not
in self.convert_map
}
)
assert not missing_func_types, f"Unsupported function types
{missing_func_types}"
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
