gemini-code-assist[bot] commented on code in PR #18384:
URL: https://github.com/apache/tvm/pull/18384#discussion_r2448623637
##########
python/tvm/meta_schedule/tune_context.py:
##########
@@ -53,9 +53,8 @@ def _normalize_mod(mod: Union[PrimFunc, IRModule]) ->
IRModule:
if not isinstance(mod, IRModule):
raise TypeError(f"Expected `mod` to be PrimFunc or IRModule, but gets:
{mod}")
func_names = mod.get_global_vars()
- (func_name,) = func_names
- if len(func_names) == 1 and func_name.name_hint != "main":
- mod = IRModule({"main": mod[func_name]})
+ if len(func_names) == 1 and func_names[0].name_hint != "main":
+ mod = IRModule({"main": mod[func_names[0]]})
Review Comment:
While the fix is correct, you can make the code slightly more concise and
readable by using an assignment expression (walrus operator `:=`). This avoids
repeating `func_names[0]` and gives it a meaningful name within the expression,
improving maintainability.
```suggestion
if len(func_names) == 1 and (gvar := func_names[0]).name_hint != "main":
mod = IRModule({"main": mod[gvar]})
```
--
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]
