style: clean up comments

Author: ayulockedin

flax/nnx/nn/attention.py

@@ -239,8 +239,8 @@ def dot_product_attention(
   query, key, value = promote_dtype((query, key, value), dtype=dtype)  # type: ignore[bad-unpacking]
   dtype = query.dtype
 
-  # GQA: Broadcast value heads to match query heads if needed.
-  # 1. Handle Key Broadcasting
+  # broadcast value heads to match query heads if needed.
+  # handle key broadcasting
   if query.ndim == key.ndim and query.shape[-2] != key.shape[-2]:
     q_heads = query.shape[-2]
     k_heads = key.shape[-2]
@@ -249,7 +249,7 @@ def dot_product_attention(
     n_rep = q_heads // k_heads
     key = jnp.repeat(key, n_rep, axis=-2)
 
-  # 2. Handle Value Broadcasting
+  # handle value broadcasting
   if query.ndim == value.ndim and query.shape[-2] != value.shape[-2]:
     q_heads = query.shape[-2]
     v_heads = value.shape[-2]

tests/nnx/nn/gqa_test.py

@@ -4,36 +4,30 @@
 
 class TestGQA:
   def test_gqa_broadcasting(self):
-    # 1. Define Shapes
     B, T, S = 2, 4, 5
     D = 8
 
     # GQA Config: Query=6 heads, Key/Value=3 heads (Ratio=2)
     num_heads_q = 6
     num_heads_kv = 3
 
-    # 2. Create Inputs
     k1, k2, k3 = jax.random.split(jax.random.key(0), 3)
     query = jax.random.normal(k1, (B, T, num_heads_q, D))
     key   = jax.random.normal(k2, (B, S, num_heads_kv, D))
     value = jax.random.normal(k3, (B, S, num_heads_kv, D))
 
-    # 3. Run Attention (Should not crash)
     output = nnx.dot_product_attention(query, key, value)
 
-    # 4. Verify Output Shape matches Query heads (6), not Key heads (3)
     assert output.shape == (B, T, num_heads_q, D)
 
   def test_gqa_invalid_heads(self):
-    # Test that it raises an error if heads aren't divisible
     B, T, D = 1, 4, 8
-    query = jnp.ones((B, T, 5, D)) # 5 heads
-    key   = jnp.ones((B, T, 2, D)) # 2 heads (5 is not divisible by 2)
+    query = jnp.ones((B, T, 5, D))
+    key   = jnp.ones((B, T, 2, D))
     value = key
 
     try:
       nnx.dot_product_attention(query, key, value)
       assert False, "Should have raised ValueError"
     except ValueError as e:
-      # Adjusted to match the actual error message in attention.py
       assert "must be multiple" in str(e)