Repository: google/flax

Issue:
NNX attention layer missing `qkv_features` arg
In Linen, the input arg [`qkv_features`](https://github.com/google/flax/blob/main/flax/linen/attention.py#L431) determines the [head dimension](https://github.com/google/flax/blob/main/flax/linen/attention.py#L436) while the input arg [`out_features`](https://github.com/google/flax/blob/main/flax/linen/attention.py#L430) determines the [output dimension](https://github.com/google/flax/blob/main/flax/linen/attention.py#L559).

In NNX, the input arg [`out_features`](https://github.com/google/flax/blob/main/flax/experimental/nnx/nnx/nn/attention.py#L318) determines **BOTH** the [head dimension](https://github.com/google/flax/blob/main/flax/experimental/nnx/nnx/nn/attention.py#L344) and the [output dimension](https://github.com/google/flax/blob/main/flax/experimental/nnx/nnx/nn/attention.py#L376).

We can get equivalent behavior with the following instantiation:
```
x = jnp.ones((12, 8))
num_heads = 4

linen_attention = nn.MultiHeadDotProductAttention(num_heads)
linen_out, variables = linen_attention.init_with_output(jax.random.key(0), x)

nnx_attention = nnx.MultiHeadAttention(num_heads, in_features=8, rngs=nnx.Rngs(0))
nnx_out = nnx_attention(x)

assert linen_out.shape == nnx_out.shape

for qkvo in ('query', 'key', 'value', 'out'):
  for kb in ('kernel', 'bias'):
    linen_shape = variables['params'][qkvo][kb].shape
    nnx_shape = getattr(getattr(nnx_attention, qkvo), kb).shape
    print(qkvo, kb, linen_shape, nnx_shape)
```
```
query kernel (8, 4, 2) (8, 4, 2)
query bias (4, 2) (4, 2)
key kernel (8, 4, 2) (8, 4, 2)
key bias (4, 2) (4, 2)
value kernel (8, 4, 2) (8, 4, 2)
value bias (4, 2) (4, 2)
out kernel (4, 2, 8) (4, 2, 8)
out bias (8,) (8,)
```
But if we add in a `qkv_features` arg input into the Linen layer:
```
linen_attention = nn.MultiHeadDotProductAttention(num_heads, qkv_features=16)
```
We'll get a different result:
```
query kernel (8, 4, 4) (8, 4, 2)
query bias (4, 4) (4, 2)
key kernel (8, 4, 4) (8, 4, 2)
key bias (4, 4) (4, 2)
value kernel (8, 4, 4) (8, 4, 2)
value bias (4, 4) (4, 2)
out kernel (4, 4, 8) (4, 2, 8)
out bias (8,) (8,)
```
Currently NNX can't replicate this behavior. Should we add a `qkv_features` argument to the NNX attention layer?


Discussion:
cc: @cgarciae 

Resolving patch:
diff --git a/flax/experimental/nnx/nnx/nn/attention.py b/flax/experimental/nnx/nnx/nn/attention.py
--- a/flax/experimental/nnx/nnx/nn/attention.py
+++ b/flax/experimental/nnx/nnx/nn/attention.py
@@ -292,6 +292,7 @@ def __init__(
     self,
     num_heads: int,
     in_features: int,
+    qkv_features: int | None = None,
     out_features: int | None = None,
     *,
     dtype: Dtype | None = None,
@@ -315,6 +316,9 @@ def __init__(
   ):
     self.num_heads = num_heads
     self.in_features = in_features
+    self.qkv_features = (
+      qkv_features if qkv_features is not None else in_features
+    )
     self.out_features = (
       out_features if out_features is not None else in_features
     )
@@ -335,13 +339,13 @@ def __init__(
     self.qkv_dot_general_cls = qkv_dot_general_cls
     self.out_dot_general_cls = out_dot_general_cls
 
-    if self.out_features % self.num_heads != 0:
+    if self.qkv_features % self.num_heads != 0:
       raise ValueError(
-        f"'out_features' ({self.out_features}) must be divisible by "
+        f'Memory dimension ({self.qkv_features}) must be divisible by '
         f"'num_heads' heads ({self.num_heads})."
       )
 
-    self.head_dim = self.out_features // self.num_heads
+    self.head_dim = self.qkv_features // self.num_heads
 
     linear_general = functools.partial(
       LinearGeneral,


Test patch:
diff --git a/flax/experimental/nnx/tests/nn/test_attention.py b/flax/experimental/nnx/tests/nn/test_attention.py
--- a/flax/experimental/nnx/tests/nn/test_attention.py
+++ b/flax/experimental/nnx/tests/nn/test_attention.py
@@ -19,7 +19,13 @@
 
 class TestMultiHeadAttention:
   def test_basic(self):
-    module = nnx.MultiHeadAttention(2, 3, 6, rngs=nnx.Rngs(0))
+    module = nnx.MultiHeadAttention(
+      num_heads=2,
+      in_features=3,
+      qkv_features=6,
+      out_features=6,
+      rngs=nnx.Rngs(0),
+    )
     y = module(jnp.ones((1, 7, 3)))
     assert y.shape == (1, 7, 6)