test fix

Author: LeiWang1999

examples/fusedmoe/example_fusedmoe_tilelang.py

@@ -7,8 +7,6 @@
 from tilelang.autotuner import *
 from example_fusedmoe_torch import *
 
-# tilelang.disable_cache()
-
 
 @tilelang.jit(pass_configs={"tl.disable_tma_lower": True, "tl.disable_warp_specialized": True})
 def moe_forward_tilelang_shared(d_hidden,

tilelang/language/proxy.py

@@ -1,7 +1,7 @@
 """The language interface for tl programs."""
 
 from __future__ import annotations
-from typing import Any, Optional, Sequence, SupportsIndex, TYPE_CHECKING, Tuple
+from typing import Any, Optional, Sequence, SupportsIndex, TYPE_CHECKING, Tuple, Union
 from typing_extensions import Self
 
 from tvm import tir
@@ -150,7 +150,9 @@ def _construct_strides(shape: Tuple[Any]):
             strides.append(s)
         return tuple(reversed(strides))
 
-    def __call__(self, shape: Tuple[Any], dtype: str = "float32", data=None) -> tir.Buffer:
+    def __call__(self, shape: Union[Tuple[Any], PrimExpr, int], dtype: str = "float32", data=None) -> tir.Buffer:
+        if isinstance(shape, (int, PrimExpr)):
+            shape = (shape,)
         return super().__call__(
             shape, dtype=dtype, strides=TensorProxy._construct_strides(shape), data=data)