[Lint]: [pre-commit.ci] auto fixes [...]

Author: pre-commit-ci[bot]

testing/python/language/test_tilelang_language_get_warp_info.py

@@ -34,9 +34,7 @@ def laneid_kernel(A: T.Tensor((num_threads,), "int32")):
 
 
 @tilelang.jit(out_idx=[-1])
-def _get_warp_idx_sync_kernel(
-    num_threads: int = 128, warp_size: Optional[int] = None
-):
+def _get_warp_idx_sync_kernel(num_threads: int = 128, warp_size: Optional[int] = None):
 
     @T.prim_func
     def warp_idx_sync_kernel(A: T.Tensor((num_threads,), "int32")):
@@ -76,9 +74,7 @@ def warp_group_idx_kernel(A: T.Tensor((num_threads,), "int32")):
 
 
 @tilelang.jit(out_idx=[-1])
-def _shuffle_elect_kernel(
-    num_threads: int = 128, thread_extent: int = 64
-):
+def _shuffle_elect_kernel(num_threads: int = 128, thread_extent: int = 64):
 
     @T.prim_func
     def shuffle_elect_kernel(A: T.Tensor((num_threads,), "int32")):
@@ -96,24 +92,18 @@ def run_get_lane_id(num_threads: int = 128, warp_size: Optional[int] = None):
     print(kernel.get_kernel_source())
     print(A)
     expected_warp_size = _resolve_warp_size(warp_size)
-    ref = torch.arange(
-        num_threads, dtype=A.dtype, device=A.device
-    ) % expected_warp_size
+    ref = torch.arange(num_threads, dtype=A.dtype, device=A.device) % expected_warp_size
     torch.testing.assert_close(A.cpu(), ref.cpu())
     return A
 
 
-def run_get_warp_idx_sync(
-    num_threads: int = 128, warp_size: Optional[int] = None
-):
+def run_get_warp_idx_sync(num_threads: int = 128, warp_size: Optional[int] = None):
     kernel = _get_warp_idx_sync_kernel(num_threads, warp_size)
     A = kernel()
     print(kernel.get_kernel_source())
     print(A)
     expected_warp_size = _resolve_warp_size(warp_size)
-    ref = torch.arange(
-        num_threads, dtype=A.dtype, device=A.device
-    ) // expected_warp_size
+    ref = torch.arange(num_threads, dtype=A.dtype, device=A.device) // expected_warp_size
     torch.testing.assert_close(A.cpu(), ref.cpu())
     return A
 
@@ -124,9 +114,7 @@ def run_get_warp_idx(num_threads: int = 128, warp_size: Optional[int] = None):
     print(kernel.get_kernel_source())
     print(A)
     expected_warp_size = _resolve_warp_size(warp_size)
-    ref = torch.arange(
-        num_threads, dtype=A.dtype, device=A.device
-    ) // expected_warp_size
+    ref = torch.arange(num_threads, dtype=A.dtype, device=A.device) // expected_warp_size
     torch.testing.assert_close(A.cpu(), ref.cpu())
     return A
 
@@ -145,25 +133,19 @@ def run_get_warp_group_idx(
     threads_per_group = expected_warp_size * expected_warps_per_group
     if threads_per_group <= 0:
         raise ValueError("threads_per_group must be positive.")
-    ref = torch.arange(
-        num_threads, dtype=A.dtype, device=A.device
-    ) // threads_per_group
+    ref = torch.arange(num_threads, dtype=A.dtype, device=A.device) // threads_per_group
     torch.testing.assert_close(A.cpu(), ref.cpu())
     return A
 
 
-def run_shuffle_elect(
-    num_threads: int = 128, thread_extent: int = 64
-):
+def run_shuffle_elect(num_threads: int = 128, thread_extent: int = 64):
     if thread_extent < 0:
         raise ValueError("thread_extent must be non-negative.")
     kernel = _shuffle_elect_kernel(num_threads, thread_extent)
     A = kernel()
     print(kernel.get_kernel_source())
     print(A)
-    indices = torch.arange(
-        num_threads, device=A.device, dtype=torch.int64
-    )
+    indices = torch.arange(num_threads, device=A.device, dtype=torch.int64)
     if thread_extent == 0:
         mask = indices == 0
     elif thread_extent > 0:
@@ -224,6 +206,7 @@ def test_shuffle_elect_default():
 def test_shuffle_elect_block_leader():
     run_shuffle_elect(num_threads=128, thread_extent=0)
 
+
 if __name__ == "__main__":
     tilelang.testing.main()
     # run_get_lane_id()