diff --git a/python/tvm/script/parser/tir/__init__.py b/python/tvm/script/parser/tir/__init__.py index 7754baf087f5..ad16821a89a3 100644 --- a/python/tvm/script/parser/tir/__init__.py +++ b/python/tvm/script/parser/tir/__init__.py @@ -16,10 +16,20 @@ # under the License. """The tir parser""" +from typing import TYPE_CHECKING + from ...ir_builder.tir import * # pylint: disable=redefined-builtin from ...ir_builder.tir import ir as _tir from . import operation as _operation from . import parser as _parser -from .entry import Buffer, Ptr, prim_func +from .entry import Buffer, Ptr + +if TYPE_CHECKING: + # pylint: disable=invalid-name + # Define prim_func and make it type check as static method + # so most tvmscript won't trigger pylint error here. + prim_func = staticmethod +else: + from .entry import prim_func __all__ = _tir.__all__ + ["Buffer", "Ptr", "prim_func"]