style: apply ruff format to test and utility files

Author: zhangyue

tests/test_flash_attention.py

@@ -307,9 +307,7 @@ def test_flash_attention_decode_cpu_cuseqlens(
         dtype=dtype,
         device=device,
     )
-    output = torch.empty(
-        (num_reqs, num_heads, head_size), dtype=dtype, device=device
-    )
+    output = torch.empty((num_reqs, num_heads, head_size), dtype=dtype, device=device)
 
     block_table = torch.zeros(
         (num_reqs, num_blocks_per_req), dtype=torch.int32, device=device
@@ -319,9 +317,7 @@ def test_flash_attention_decode_cpu_cuseqlens(
         for j in range(num_blocks_per_req):
             block_table[i, j] = i * num_blocks_per_req + j
 
-    cu_seqlens_q = torch.arange(
-        0, num_reqs + 1, dtype=torch.int64, device=device
-    )
+    cu_seqlens_q = torch.arange(0, num_reqs + 1, dtype=torch.int64, device=device)
 
     # CPU cu_seqlens_kv — exercises `detail::extractSeqLengths` host path
     # (direct pointer read, no D2H copy).

tests/test_paged_attention.py

@@ -384,9 +384,7 @@ def test_paged_attention_host_tensors(
         dtype=dtype,
         device=device,
     )
-    output = torch.empty(
-        (num_reqs, num_heads, head_size), dtype=dtype, device=device
-    )
+    output = torch.empty((num_reqs, num_heads, head_size), dtype=dtype, device=device)
 
     block_table = torch.zeros(
         (num_reqs, num_blocks_per_req), dtype=torch.int32, device=device
@@ -396,9 +394,7 @@ def test_paged_attention_host_tensors(
         for j in range(num_blocks_per_req):
             block_table[i, j] = i * num_blocks_per_req + j
 
-    seq_lens = torch.full(
-        (num_reqs,), kv_len, dtype=torch.int32, device=device
-    )
+    seq_lens = torch.full((num_reqs,), kv_len, dtype=torch.int32, device=device)
 
     # CPU copies for the D2H-free path.
     seq_lens_cpu = seq_lens.cpu().contiguous()

tests/test_reshape_and_cache.py

@@ -55,9 +55,7 @@ def test_reshape_and_cache_contiguous(
     active_indices = infini.ops.ReshapeAndCache.active_implementation_indices(device)
 
     if implementation_index not in active_indices:
-        pytest.skip(
-            f"implementation `{implementation_index}` not active on `{device}`"
-        )
+        pytest.skip(f"implementation `{implementation_index}` not active on `{device}`")
 
     key = randn_strided(
         (num_tokens, num_kv_heads, head_size), None, dtype=dtype, device=device
@@ -125,9 +123,7 @@ def test_reshape_and_cache_noncontiguous_slots(
     active_indices = infini.ops.ReshapeAndCache.active_implementation_indices(device)
 
     if implementation_index not in active_indices:
-        pytest.skip(
-            f"implementation `{implementation_index}` not active on `{device}`"
-        )
+        pytest.skip(f"implementation `{implementation_index}` not active on `{device}`")
 
     key = randn_strided(
         (num_tokens, num_kv_heads, head_size), None, dtype=dtype, device=device
@@ -157,17 +153,26 @@ def test_reshape_and_cache_noncontiguous_slots(
     )
 
 
-def _reshape_and_cache(key, value, kv_cache, slot_mapping, kv_cache_out,
-                       implementation_index=0):
+def _reshape_and_cache(
+    key, value, kv_cache, slot_mapping, kv_cache_out, implementation_index=0
+):
     if key.device.type == "npu":
         infini.ops.reshape_and_cache(
-            key, value, kv_cache, slot_mapping, kv_cache_out,
+            key,
+            value,
+            kv_cache,
+            slot_mapping,
+            kv_cache_out,
             implementation_index=implementation_index,
             stream=get_npu_stream(key),
         )
     else:
         infini.ops.reshape_and_cache(
-            key, value, kv_cache, slot_mapping, kv_cache_out,
+            key,
+            value,
+            kv_cache,
+            slot_mapping,
+            kv_cache_out,
             implementation_index=implementation_index,
         )
 

tests/utils.py

@@ -100,5 +100,3 @@ def clone_strided(input):
     output.as_strided(*as_strided_args).copy_(input.as_strided(*as_strided_args))
 
     return output
-
-