Fix lint

Author: BestJuly

tests/unit_tests/ssm/test_gated_delta_net.py

@@ -149,29 +149,36 @@ def test_jit_compiled_helpers(self):
 
         qkv_last_dim = (2 * gdn.qk_dim_local_tp + gdn.v_dim_local_tp) // gdn.cp_size
         qkv = torch.randn(
-            batch, seq_len, qkv_last_dim,
-            device=torch.cuda.current_device(), dtype=torch.bfloat16,
+            batch, seq_len, qkv_last_dim, device=torch.cuda.current_device(), dtype=torch.bfloat16
         )
         gate = torch.randn(
-            batch, seq_len, num_v_heads_local, gdn.value_head_dim,
-            device=torch.cuda.current_device(), dtype=torch.bfloat16,
+            batch,
+            seq_len,
+            num_v_heads_local,
+            gdn.value_head_dim,
+            device=torch.cuda.current_device(),
+            dtype=torch.bfloat16,
         )
         beta = torch.randn(
-            batch, seq_len, num_v_heads_local,
-            device=torch.cuda.current_device(), dtype=torch.bfloat16,
+            batch,
+            seq_len,
+            num_v_heads_local,
+            device=torch.cuda.current_device(),
+            dtype=torch.bfloat16,
         )
         alpha = torch.randn(
-            batch, seq_len, num_v_heads_local,
-            device=torch.cuda.current_device(), dtype=torch.bfloat16,
+            batch,
+            seq_len,
+            num_v_heads_local,
+            device=torch.cuda.current_device(),
+            dtype=torch.bfloat16,
         )
 
         # Disable dynamo so coverage.py can trace through the method bodies,
         # which are normally wrapped by @jit_fuser (torch.compile).
         with torch._dynamo.config.patch(disable=True):
             query, key, value, gate_out, beta_out, alpha_out = (
-                gdn._prepare_qkv_for_gated_delta_rule(
-                    qkv, gate, beta, alpha, batch, seq_len
-                )
+                gdn._prepare_qkv_for_gated_delta_rule(qkv, gate, beta, alpha, batch, seq_len)
             )
 
         assert query.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim)
@@ -182,16 +189,14 @@ def test_jit_compiled_helpers(self):
         assert value.is_contiguous()
 
         A_log_mock = torch.randn(
-            num_v_heads_local, device=torch.cuda.current_device(), dtype=torch.bfloat16,
+            num_v_heads_local, device=torch.cuda.current_device(), dtype=torch.bfloat16
         )
         dt_bias_mock = torch.randn(
-            num_v_heads_local, device=torch.cuda.current_device(), dtype=torch.bfloat16,
+            num_v_heads_local, device=torch.cuda.current_device(), dtype=torch.bfloat16
         )
 
         with torch._dynamo.config.patch(disable=True):
-            g, beta_sig = gdn._compute_g_and_beta(
-                A_log_mock, dt_bias_mock, alpha, beta
-            )
+            g, beta_sig = gdn._compute_g_and_beta(A_log_mock, dt_bias_mock, alpha, beta)
 
         assert g.dtype == torch.float32
         assert g.shape == alpha.shape