lint

Author: xwhzz

testing/python/jit/test_tilelang_jit_nullptr.py

@@ -5,6 +5,7 @@
 import tilelang.language as T
 from tilelang.utils import map_torch_type
 
+
 @tl.jit
 def ptr_null_test(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float"):
 
@@ -46,16 +47,17 @@ def main(
 
     return main
 
+
 @tl.jit
 def tensor_null_test(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float"):
 
     @T.prim_func
     def main(
-        A: T.Tensor((M, K), dtype),
-        B: T.Tensor((K, N), dtype),
-        C: T.Tensor((M, N), accum_dtype),
-        Bias: T.Tensor((N), accum_dtype),
-        with_bias: T.bool,
+            A: T.Tensor((M, K), dtype),
+            B: T.Tensor((K, N), dtype),
+            C: T.Tensor((M, N), accum_dtype),
+            Bias: T.Tensor((N), accum_dtype),
+            with_bias: T.bool,
     ):
         # Initialize Kernel Context
         with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
@@ -79,6 +81,7 @@ def main(
 
     return main
 
+
 def run_test(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float"):
     func = ptr_null_test(M, N, K, block_M, block_N, block_K, dtype, accum_dtype)
 
@@ -104,8 +107,10 @@ def run_test(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="f
     func(a, b, c, d, True)
     torch.testing.assert_close(c, ref_with_bias, atol=1e-2, rtol=1e-2)
 
+
 def test_nullptr():
     run_test(1024, 1024, 1024, 128, 128, 32)
 
+
 if __name__ == "__main__":
-    tilelang.testing.main()
+    tilelang.testing.main()