fix some remaining import issue

Author: oraluben

tilelang/cache/__init__.py

@@ -1,7 +1,7 @@
 """The cache utils with class and database persistence - Init file"""
 from __future__ import annotations
 
-from typing import List, Union, Literal, Optional
+from typing import Literal
 from tvm.target import Target
 from tvm.tir import PrimFunc
 from tilelang.jit import JITKernel

tilelang/carver/__init__.py

@@ -1,26 +1,10 @@
 """Base infra"""
 
-from .analysis import (
-    BlockInfo,  # noqa: F401
-    IterInfo,  # noqa: F401
-    collect_block_iter_vars_used_in_access_region,  # noqa: F401
-    collect_vars_used_in_prim_expr,  # noqa: F401
-    detect_dominant_read,  # noqa: F401
-    is_broadcast_epilogue,  # noqa: F401
-    normalize_prim_func,  # noqa: F401
-)  # noqa: F401
-from .common_schedules import (
-    get_block,
-    get_output_blocks,
-    try_inline,
-    try_inline_contiguous_spatial,
-)  # noqa: F401
+from .analysis import *
+from .common_schedules import *
 from .roller import *
-from .arch import CUDA, CDNA  # noqa: F401
-from .template import (
-    MatmulTemplate,
-    GEMVTemplate,
-    ElementwiseTemplate,
-    GeneralReductionTemplate,
-    FlashAttentionTemplate,
-)  # noqa: F401
+from .arch import (
+    CUDA as CUDA,
+    CDNA as CDNA,
+)
+from .template import *

tilelang/carver/analysis.py

@@ -297,3 +297,14 @@ def get_coalesced_veclen(block_stmt: tir.Block, target_bits: int = 128) -> int:
     for buffer in buffers:
         max_dtype_bits = max(max_dtype_bits, DataType(buffer.dtype).bits)
     return target_bits // max_dtype_bits
+
+
+__all__ = [
+    BlockInfo,
+    IterInfo,
+    collect_block_iter_vars_used_in_access_region,
+    collect_vars_used_in_prim_expr,
+    detect_dominant_read,
+    is_broadcast_epilogue,
+    normalize_prim_func,
+]

tilelang/carver/arch/__init__.py

@@ -4,7 +4,6 @@
 from .cpu import *
 from .cdna import *
 from .metal import *
-from typing import Union
 from tvm.target import Target
 import torch
 

tilelang/carver/common_schedules.py

@@ -163,3 +163,11 @@ def try_inline_contiguous_spatial(
     if spatial_blocks:
         results.extend(try_inline(sch, spatial_blocks))
     return results
+
+
+__all__ = [
+    get_block,
+    get_output_blocks,
+    try_inline,
+    try_inline_contiguous_spatial,
+]

tilelang/carver/template/__init__.py

@@ -1,9 +1,19 @@
 """Template for the TileLang Carver."""
 
-from .base import BaseTemplate  # noqa: F401
-from .matmul import MatmulTemplate  # noqa: F401
-from .gemv import GEMVTemplate  # noqa: F401
-from .elementwise import ElementwiseTemplate  # noqa: F401
-from .general_reduce import GeneralReductionTemplate  # noqa: F401
-from .flashattention import FlashAttentionTemplate  # noqa: F401
-from .conv import ConvTemplate  # noqa: F401
+from .base import BaseTemplate
+from .matmul import MatmulTemplate
+from .gemv import GEMVTemplate
+from .elementwise import ElementwiseTemplate
+from .general_reduce import GeneralReductionTemplate
+from .flashattention import FlashAttentionTemplate
+from .conv import ConvTemplate
+
+__all__ = [
+    'BaseTemplate',
+    'MatmulTemplate',
+    'GEMVTemplate',
+    'ElementwiseTemplate',
+    'GeneralReductionTemplate',
+    'FlashAttentionTemplate',
+    'ConvTemplate',
+]

tilelang/jit/__init__.py

@@ -7,14 +7,9 @@
 
 from typing import (
     Any,
-    List,
-    Union,
     Callable,
-    Tuple,
     overload,
     Literal,
-    Dict,  # For type hinting dicts
-    Optional,
 )
 from tilelang import tvm as tvm
 from tilelang.jit.adapter.utils import is_metal_target

tilelang/language/__init__.py

@@ -1,7 +1,7 @@
 """The language interface for tl programs."""
 from __future__ import annotations
 
-from typing import Optional, Callable, Dict
+from typing import Callable
 
 # from .parser import *
 # now is fully compatible with the upstream

tilelang/primitives/gemm/__init__.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Optional
+
 from tvm import tir
 from tilelang.utils import is_local, is_fragment, is_shared
 from tilelang.primitives.gemm.base import GemmWarpPolicy

tilelang/profiler/__init__.py

@@ -1,7 +1,7 @@
 """The profiler and convert to torch utils"""
 from __future__ import annotations
 
-from typing import List, Optional, Callable, Any, Literal
+from typing import Callable, Any, Literal
 from functools import partial
 import torch
 from contextlib import suppress