Silence some pytype errors.

PiperOrigin-RevId: 623308993

Author: rchen152

jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py

@@ -418,9 +418,9 @@ def _wrapped(
     if is_grouped:
 
       def reshape_activations(activations):
-        if activations.ndim == 4:
-          kv_heads, q_heads_per_kv_head, q_seq_len, head_dim = activations.shape
-          return activations.reshape(
+        if activations.ndim == 4:  # pytype: disable=attribute-error
+          kv_heads, q_heads_per_kv_head, q_seq_len, head_dim = activations.shape  # pytype: disable=attribute-error
+          return activations.reshape(  # pytype: disable=attribute-error
               kv_heads * q_heads_per_kv_head, q_seq_len, head_dim
           )
         return activations
@@ -755,7 +755,7 @@ def body(kv_compute_index, _):
 
     qk = apply_mask_and_soft_cap()
 
-    m_curr = qk.max(axis=-1)[:, None]
+    m_curr = qk.max(axis=-1)[:, None]  # pytype: disable=attribute-error
     assert m_curr.shape == (bq, 1)
     m_next = jnp.maximum(m_prev, m_curr)
     assert m_next.shape == (bq, NUM_LANES)