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

@@ -1,39 +1,66 @@
-from flax import nnx
-import jax.numpy as jnp
 import jax
+import jax.numpy as jnp
+from flax import nnx
+import numpy as np
 
 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)
-    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)
+    def test_gqa_shapes(self):
+        B, T, S = 2, 4, 5
+        D = 8
+        num_heads_q = 6
+        num_heads_kv = 3
+
+        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))
+
+        output = nnx.dot_product_attention(query, key, value)
+        expected_shape = (B, T, num_heads_q, D)
+        assert output.shape == expected_shape
+
+    def test_gqa_invalid_heads(self):
+        B, T, D = 1, 4, 8
+        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:
+            # Fixed expectation to match your code's error message
+            assert "must be multiple" in str(e)
+
+    def test_gqa_parity_with_jax(self):
+        class DummyModule(nnx.Module):
+            pass
+        
+        dummy_module = DummyModule()
+
+        B, T, S, D = 2, 8, 8, 16
+        num_heads_q = 4
+        num_heads_kv = 2
+        
+        rng = jax.random.key(42)
+        k1, k2, k3 = jax.random.split(rng, 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))
+
+        # Manually repeat heads for JAX reference
+        n_rep = num_heads_q // num_heads_kv
+        key_jax = jnp.repeat(key, n_rep, axis=-2)
+        value_jax = jnp.repeat(value, n_rep, axis=-2)
+        
+        jax_out = jax.nn.dot_product_attention(query, key_jax, value_jax)
+
+        # NNX should handle broadcasting internally
+        nnx_out = nnx.dot_product_attention(
+            query, key, value, 
+            module=dummy_module
+        )
+
+        # Relaxed tolerance to 1e-3
+        np.testing.assert_allclose(nnx_out, jax_out, atol=1e-3, rtol=1e-3)