gemini-code-assist[bot] commented on code in PR #19630: URL: https://github.com/apache/tvm/pull/19630#discussion_r3314262522
########## python/tvm/tirx/functor.py: ########## @@ -23,7 +23,6 @@ import tvm_ffi from tvm.ir import PrimExpr -from tvm.runtime.support import derived_object from . import _ffi_api Review Comment:  Import `TYPE_CHECKING` from `typing` to support static analysis and IDE autocompletion for the lazily loaded attributes. ```python from typing import TYPE_CHECKING from tvm.ir import PrimExpr ``` ########## python/tvm/tirx/functor.py: ########## @@ -78,39 +77,26 @@ While, ) -visitor = derived_object -""" -A decorator to wrap user-customized PyStmtExprVisitor as TVM object _PyStmtExprVisitor. -Parameters ----------- -visitor_cls : PyStmtExprVisitor - The user-customized PyStmtExprVisitor. +def __getattr__(name: str): + # ``derived_object``, ``visitor``, and ``mutator`` are loaded lazily to avoid + # a circular import: tvm.tirx is imported before tvm.s_tir is fully + # initialized, so a top-level ``from tvm.s_tir.utils import derived_object`` + # would trigger tvm.s_tir.__init__ mid-flight and hit a circular dependency. + if name in ("derived_object", "visitor", "mutator"): + from tvm.s_tir.utils import derived_object # pylint: disable=import-outside-toplevel -Returns -------- -cls : _PyStmtExprVisitor - The decorated TVM object _PyStmtExprVisitor(StmtExprVisitor on the C++ side). + globals()["derived_object"] = derived_object + globals()["visitor"] = derived_object + globals()["mutator"] = derived_object + return derived_object + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") -Example -------- -.. code-block:: python - @tirx.functor.stmt_expr_visitor - class MyStmtExprVisitor(PyStmtExprVisitor): - # customize visit function - def visit_call_(self, op: Call) -> None: - # just for demo purposes - ... - # myvisitor is now a special visitor that visit every Call with - # user-customized visit_call_ - myvisitor = MyStmtExprVisitor() - # apply myvisitor to PrimExpr and Stmt - myvisitor.visit_expr(expr) - myvisitor.visit_stmt(stmt) -""" +# visitor and mutator are aliases for derived_object, loaded lazily via module __getattr__ +# to avoid a circular import at package init time. See __getattr__ defined above. + Review Comment:  Add a `if TYPE_CHECKING:` block to declare `visitor` and `mutator` at the module level for static analysis tools (like mypy, pylint) and IDE autocompletion. This also allows us to preserve the docstring for `visitor` which was otherwise lost during the refactoring. ```python def __getattr__(name: str): # ``derived_object``, ``visitor``, and ``mutator`` are loaded lazily to avoid # a circular import: tvm.tirx is imported before tvm.s_tir is fully # initialized, so a top-level ``from tvm.s_tir.utils import derived_object`` # would trigger tvm.s_tir.__init__ mid-flight and hit a circular dependency. if name in ("derived_object", "visitor", "mutator"): from tvm.s_tir.utils import derived_object # pylint: disable=import-outside-toplevel globals()["derived_object"] = derived_object globals()["visitor"] = derived_object globals()["mutator"] = derived_object return derived_object raise AttributeError(f"module {__name__!r} has no attribute {name!r}") if TYPE_CHECKING: from tvm.s_tir.utils import derived_object visitor = derived_object """A decorator to wrap user-customized PyStmtExprVisitor as TVM object _PyStmtExprVisitor. Parameters ---------- visitor_cls : PyStmtExprVisitor The user-customized PyStmtExprVisitor. Returns ------- cls : _PyStmtExprVisitor The decorated TVM object _PyStmtExprVisitor(StmtExprVisitor on the C++ side). """ mutator = derived_object ``` -- 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]
