Merge pull request jax-ml#27258 from jakevdp:fix-lint

Google-ML-Automation · Google-ML-Automation · commit cf21f733d439 · 2025-03-19T12:13:12.000-07:00
PiperOrigin-RevId: 738492394
diff --git a/jax/experimental/pallas/ops/tpu/ragged_paged_attention.py b/jax/experimental/pallas/ops/tpu/ragged_paged_attention.py
@@ -431,7 +431,7 @@ def init_scratch_ref():
       causal_mask = row_ids < col_ids
       if sliding_window is not None:
         causal_mask = jnp.logical_or(causal_mask,
-                                     row_ids - sliding_window>=col_ids)
+                                     row_ids - sliding_window >= col_ids)
       qk += jnp.where(causal_mask, mask_value, 0.0)
       m_curr = jnp.max(qk, axis=1, keepdims=True)
       s_curr = jnp.exp(qk - m_curr)
diff --git a/tests/pallas/tpu_ragged_paged_attention_test.py b/tests/pallas/tpu_ragged_paged_attention_test.py
@@ -306,7 +306,7 @@ def test_ragged_paged_attention_complex(
     )
 
   def test_ragged_paged_attention_sliding_window_should_be_positive(self):
-    dtype=jnp.float32
+    dtype = jnp.float32
     seq_lens = [(192, 328), (128, 180), (64, 255)]
     num_heads = (32, 8)
     head_dim = 128

Original file line number	Diff line number	Diff line change
`@@ -306,7 +306,7 @@ def test_ragged_paged_attention_complex(`
`306`	`306`	`)`
`307`	`307`
`308`	`308`	`def test_ragged_paged_attention_sliding_window_should_be_positive(self):`
`309`		`- dtype=jnp.float32`
	`309`	`+ dtype = jnp.float32`
`310`	`310`	`seq_lens = [(192, 328), (128, 180), (64, 255)]`
`311`	`311`	`num_heads = (32, 8)`
`312`	`312`	`head_dim = 128`