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:

Patch:
diff --git a/flax/experimental/nnx/nnx/graph_utils.py b/flax/experimental/nnx/nnx/graph_utils.py
index 3a0f3111..501bb097 100644
--- a/flax/experimental/nnx/nnx/graph_utils.py
+++ b/flax/experimental/nnx/nnx/graph_utils.py
@@ -171,6 +171,17 @@ class _HashableMapping(tp.Mapping[HA, HB], tp.Hashable):
     return repr(self._mapping)
 
 
+@dataclasses.dataclass(repr=False)
+class _SubgraphRepr(reprlib.Representable):
+  subgraphs: tp.Mapping[str, tp.Union['GraphDef[tp.Any]', int]]
+
+  def __nnx_repr__(self):
+    yield reprlib.Object(type='', value_sep=', ')
+
+    for name, subgraph in self.subgraphs.items():
+      yield reprlib.Attr(repr(name), subgraph, start='(', end=')')
+
+
 @dataclasses.dataclass(repr=False)
 class _MappingRepr(reprlib.Representable):
   mapping: tp.Mapping[str, tp.Any]
@@ -436,7 +447,7 @@ def _graph_flatten(
 
 def graph_unflatten(graphdef: GraphDef[Node], state: State) -> Node:
   index_to_node: dict[Index, tp.Any] = {}
-  return _graph_unflatten(graphdef, state.raw_mapping, index_to_node)
+  return _graph_unflatten(graphdef, state.variables, index_to_node)
 
 
 def _graph_unflatten(
@@ -610,7 +621,7 @@ def graph_update_dynamic(
     new_states = updates
 
   for state in new_states:
-    _graph_update_dynamic(node, state.raw_mapping)
+    _graph_update_dynamic(node, state.variables)
 
 
 def _graph_update_dynamic(
diff --git a/flax/experimental/nnx/nnx/module.py b/flax/experimental/nnx/nnx/module.py
index 38e37535..75156515 100644
--- a/flax/experimental/nnx/nnx/module.py
+++ b/flax/experimental/nnx/nnx/module.py
@@ -65,7 +65,7 @@ class _HasSetup(tp.Protocol):
 SEEN_MODULES_REPR: tp.Optional[tp.Set[ids.UUID]] = None
 
 
-class ModuleVariablesMapping(
+class VariablesMapping(
   tp.MutableMapping[str, Variable[tp.Any]], reprlib.Representable
 ):
   __slots__ = ('_module',)
@@ -76,17 +76,14 @@ class ModuleVariablesMapping(
     else:
       object.__setattr__(self, '_module', module)
 
-  def __getitem__(self, key: str | int) -> Variable[tp.Any]:
-    if isinstance(key, int):
-      key = str(key)
-
+  def __getitem__(self, name: str) -> Variable[tp.Any]:
     module_vars = vars(self._module)
-    if key not in module_vars:
-      raise KeyError(f'Variable {key} not found')
-    value = module_vars[key]
+    if name not in module_vars:
+      raise KeyError(f'Variable {name} not found')
+    value = module_vars[name]
 
     if not isinstance(value, Variable):
-      raise KeyError(f"Variable '{key}' is not found.")
+      raise KeyError(f"Variable '{name}' is not found.")
 
     return value
 
@@ -120,7 +117,7 @@ class ModuleVariablesMapping(
     yield reprlib.Object(type(self), start='{', end='}', value_sep=': ')
     for name, value in vars(self._module).items():
       if isinstance(value, Variable):
-        yield reprlib.Attr(repr(name), value)
+        yield reprlib.Attr(name, value)
 
 
 class ModuleState(reprlib.Representable):
@@ -238,8 +235,8 @@ class Module(reprlib.Representable, metaclass=ModuleMeta):
       vars_dict[name] = value
 
   @property
-  def variables(self) -> ModuleVariablesMapping:
-    return ModuleVariablesMapping(self)
+  def variables(self) -> VariablesMapping:
+    return VariablesMapping(self)
 
   def __deepcopy__(self: M, memo=None) -> M:
     state, graphdef = self.split()
@@ -534,7 +531,7 @@ class Module(reprlib.Representable, metaclass=ModuleMeta):
 # -------------------------
 def _module_flatten(module: Module, *, with_keys: bool):
   state, graphdef = module.split()
-  variables = state.raw_mapping
+  variables = state.variables
   paths = tuple(variables.keys())
 
   if with_keys:
diff --git a/flax/experimental/nnx/nnx/nn/attention.py b/flax/experimental/nnx/nnx/nn/attention.py
index 871b5c55..3f46da38 100644
--- a/flax/experimental/nnx/nnx/nn/attention.py
+++ b/flax/experimental/nnx/nnx/nn/attention.py
@@ -292,6 +292,7 @@ class MultiHeadAttention(Module):
     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 @@ class MultiHeadAttention(Module):
   ):
     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 @@ class MultiHeadAttention(Module):
     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,
diff --git a/flax/experimental/nnx/nnx/state.py b/flax/experimental/nnx/nnx/state.py
index 5789c54f..3b650638 100644
--- a/flax/experimental/nnx/nnx/state.py
+++ b/flax/experimental/nnx/nnx/state.py
@@ -44,58 +44,6 @@ Key = str
 FlatState = dict[Path, Variable[Leaf]]
 
 
-class StateVariablesMapping(
-  tp.MutableMapping[str, Variable[tp.Any]], reprlib.Representable
-):
-  __slots__ = ('_mapping',)
-
-  def __init__(self, mapping: dict[Key, tp.Any]):
-    if tp.TYPE_CHECKING:
-      self._mapping = mapping
-    else:
-      object.__setattr__(self, '_mapping', mapping)
-
-  def __getitem__(self, key: str | int) -> Variable[tp.Any]:
-    if isinstance(key, int):
-      key = str(key)
-
-    value = self._mapping[key]
-
-    if not isinstance(value, Variable):
-      raise KeyError(f"Variable '{key}' not found.")
-
-    return value
-
-  def __setitem__(self, name: str, value: Variable[tp.Any]) -> None:
-    self._mapping[name] = value
-
-  def __getattr__(self, name: str) -> Variable[tp.Any]:
-    value = self._mapping[name]
-    if not isinstance(value, Variable):
-      raise AttributeError(f"Variable '{name}' not found.")
-    return value
-
-  def __setattr__(self, name: str, value: Variable[tp.Any]) -> None:
-    self._mapping[name] = value
-
-  def __delitem__(self, name: str) -> None:
-    del self._mapping[name]
-
-  def __iter__(self) -> tp.Iterator[str]:
-    for name, value in self._mapping.items():
-      if isinstance(value, Variable):
-        yield name
-
-  def __len__(self) -> int:
-    return sum(1 for _ in self)
-
-  def __nnx_repr__(self):
-    yield reprlib.Object(type(self), start='{', end='}', value_sep=': ')
-    for name, value in vars(self._mapping).items():
-      if isinstance(value, Variable):
-        yield reprlib.Attr(repr(name), value)
-
-
 class NestedStateRepr(reprlib.Representable):
   def __init__(self, state: State):
     self.state = state
@@ -124,13 +72,9 @@ class State(tp.MutableMapping[Key, tp.Any], reprlib.Representable):
       super().__setattr__('_mapping', dict(mapping))
 
   @property
-  def raw_mapping(self) -> dict[Key, dict[str, tp.Any] | tp.Any]:
+  def variables(self) -> dict[Key, dict[str, tp.Any] | tp.Any]:
     return self._mapping
 
-  @property
-  def variables(self) -> StateVariablesMapping:
-    return StateVariablesMapping(self._mapping)
-
   def __getitem__(self, key: Key | int) -> Leaf | State:
     if isinstance(key, int):
       key = str(key)


Test patch:
diff --git a/flax/experimental/nnx/tests/nn/test_attention.py b/flax/experimental/nnx/tests/nn/test_attention.py
index b7be9270..bbc36f35 100644
--- a/flax/experimental/nnx/tests/nn/test_attention.py
+++ b/flax/experimental/nnx/tests/nn/test_attention.py
@@ -19,7 +19,13 @@ from flax.experimental import nnx
 
 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)
 
diff --git a/flax/experimental/nnx/tests/test_graph_utils.py b/flax/experimental/nnx/tests/test_graph_utils.py
index bf1d694b..614eac85 100644
--- a/flax/experimental/nnx/tests/test_graph_utils.py
+++ b/flax/experimental/nnx/tests/test_graph_utils.py
@@ -190,15 +190,15 @@ class TestGraphUtils:
     state, static = m.split()
 
     assert isinstance(m.variables.a, nnx.Param)
-    assert isinstance(state.variables.a, nnx.Param)
-    assert m.variables.a is not state.variables.a
+    assert isinstance(state.variables['a'], nnx.Param)
+    assert m.variables.a is not state.variables['a']
     assert m.a == state.a
 
     m2 = static.merge(state)
 
     assert isinstance(m2.variables.a, nnx.Param)
-    assert isinstance(state.variables.a, nnx.Param)
-    assert m2.variables.a is not state.variables.a
+    assert isinstance(state.variables['a'], nnx.Param)
+    assert m2.variables.a is not state.variables['a']
     assert m2.a == state.a
 
   def test_shared_state_variables_not_shared_with_graph(self):
@@ -213,10 +213,10 @@ class TestGraphUtils:
 
     assert isinstance(m.variables.a, nnx.Param)
     assert isinstance(m.variables.b, nnx.Param)
-    assert isinstance(state.variables.a, nnx.Param)
+    assert isinstance(state.variables['a'], nnx.Param)
     assert 'b' not in state
-    assert m.variables.a is not state.variables.a
-    assert m.variables.b is not state.variables.a
+    assert m.variables.a is not state.variables['a']
+    assert m.variables.b is not state.variables['a']
     assert m.a == state.a
     assert m.b == state.a
 
@@ -224,9 +224,9 @@ class TestGraphUtils:
 
     assert isinstance(m2.variables.a, nnx.Param)
     assert isinstance(m2.variables.b, nnx.Param)
-    assert isinstance(state.variables.a, nnx.Param)
-    assert m2.variables.a is not state.variables.a
-    assert m2.variables.b is not state.variables.a
+    assert isinstance(state.variables['a'], nnx.Param)
+    assert m2.variables.a is not state.variables['a']
+    assert m2.variables.b is not state.variables['a']
     assert m2.a == state.a
     assert m2.b == state.a
     assert m2.a is m2.b
diff --git a/flax/experimental/nnx/tests/test_module.py b/flax/experimental/nnx/tests/test_module.py
index 8d2a5ac8..65c537d0 100644
--- a/flax/experimental/nnx/tests/test_module.py
+++ b/flax/experimental/nnx/tests/test_module.py
@@ -281,7 +281,7 @@ class TestModule:
 
     intermediates = m.pop(nnx.Intermediate)
 
-    assert isinstance(intermediates.variables.y, nnx.Intermediate)
+    assert isinstance(intermediates.variables['y'], nnx.Intermediate)
     assert intermediates['y'] == (3, 11)
 
     assert hasattr(m, 'y')
@@ -515,10 +515,10 @@ class TestModuleDataclass:
     state, graphdef = m.split()
 
     assert len(state) == 4
-    assert state.variables.b == nnx.TreeNode(2)
-    assert state.variables.c == nnx.Param(3)
-    assert state.variables.d == nnx.BatchStat(4)
-    assert state.variables.f == nnx.Variable(6)
+    assert state.variables['b'] == nnx.TreeNode(2)
+    assert state.variables['c'] == nnx.Param(3)
+    assert state.variables['d'] == nnx.BatchStat(4)
+    assert state.variables['f'] == nnx.Variable(6)
 
   def test_no_override(self):
     @nnx.dataclass
@@ -584,7 +584,7 @@ class TestModuleDef:
     states, graphdef = foo.split()
 
     assert isinstance(states, nnx.State)
-    assert isinstance(states.variables.w, nnx.Param)
+    assert isinstance(states.variables['w'], nnx.Param)
     # assert isinstance(states["c"], jax.Array)
 
     y, _updates = graphdef.apply(states)(x=2.0, rngs=nnx.Rngs(e=1))
@@ -609,8 +609,8 @@ class TestModuleDef:
 
     assert isinstance(graphdef, nnx.GraphDef)
     assert isinstance(state, nnx.State)
-    assert isinstance(state.variables.w, nnx.Param)
-    assert isinstance(state.variables.c, nnx.Variable)
+    assert isinstance(state.variables['w'], nnx.Param)
+    assert isinstance(state.variables['c'], nnx.Variable)
 
     y, (state, graphdef) = graphdef.apply(state)(x=2.0, rngs=nnx.Rngs(e=1))
 
diff --git a/flax/experimental/nnx/tests/test_partitioning.py b/flax/experimental/nnx/tests/test_partitioning.py
index 7b9f074d..09a8b514 100644
--- a/flax/experimental/nnx/tests/test_partitioning.py
+++ b/flax/experimental/nnx/tests/test_partitioning.py
@@ -155,5 +155,5 @@ class TestPartitioning:
     assert state['a']['0'] == m.a[0]
     assert state['a']['1'] == m.a[1]
     assert state['b'] == m.b
-    assert state.variables.b is not state.a.variables[0]
+    assert state.variables['b'] is not state.variables['a']['0']
     assert len(state.flat_state()) == 3
diff --git a/flax/experimental/nnx/tests/test_state.py b/flax/experimental/nnx/tests/test_state.py
index b27da4db..2529a970 100644
--- a/flax/experimental/nnx/tests/test_state.py
+++ b/flax/experimental/nnx/tests/test_state.py
@@ -45,10 +45,10 @@ class StateTest(TestCase):
     state.a = 3
     state.b.c = 4
 
-    assert isinstance(state.variables.a, nnx.Param)
-    assert state.variables.a.value == 3
-    assert isinstance(state.b.variables.c, nnx.Param)
-    assert state.b.variables.c.value == 4
+    assert isinstance(state.variables['a'], nnx.Param)
+    assert state.variables['a'].value == 3
+    assert isinstance(state.variables['b']['c'], nnx.Param)
+    assert state.variables['b']['c'].value == 4
 
   def test_integer_access(self):
     class Foo(nnx.Module):
diff --git a/flax/experimental/nnx/tests/test_transforms.py b/flax/experimental/nnx/tests/test_transforms.py
index 8a3bab79..3ebdc8b6 100644
--- a/flax/experimental/nnx/tests/test_transforms.py
+++ b/flax/experimental/nnx/tests/test_transforms.py
@@ -139,9 +139,9 @@ class TestGrad:
     assert m.a.variables['0'] is m.variables.b
     assert isinstance(grads, nnx.State)
     assert grads['a']['0'] == 2.0
-    assert isinstance(grads.a.variables['0'], nnx.Variable)
+    assert isinstance(grads.variables['a']['0'], nnx.Variable)
     assert grads['a']['1'] == 1.0
-    assert isinstance(grads.a.variables['1'], nnx.Variable)
+    assert isinstance(grads.variables['a']['1'], nnx.Variable)
     assert len(grads.flat_state()) == 2
 
     m.update(grads)
@@ -170,7 +170,7 @@ class TestGrad:
 
     assert isinstance(grads, nnx.State)
     assert grads['a']['0'] == 1.0
-    assert isinstance(grads.a.variables['0'], nnx.Param)
+    assert isinstance(grads.variables['a']['0'], nnx.Param)
     assert len(grads) == 2
 
     m.update(grads)
@@ -198,7 +198,7 @@ class TestGrad:
 
     assert isinstance(grads, nnx.State)
     assert grads['a']['1'] == 1.0
-    assert isinstance(grads.a.variables['1'], nnx.BatchStat)
+    assert isinstance(grads.variables['a']['1'], nnx.BatchStat)
     assert len(grads) == 1
 
     m.update(grads)
@@ -463,11 +463,11 @@ class TestScan:
         x = self.linear(x)
 
         # test sharding layer axes is not present inside scan
-        state = self.linear.get_state()
-        assert state.kernel.shape == (3, 3)
-        assert state.variables.kernel.sharding == ('din', 'dout')
-        assert state.bias.shape == (3,)
-        assert state.variables.bias.sharding == ('dout',)
+        variables = self.linear.get_state().variables
+        assert variables['kernel'].value.shape == (3, 3)
+        assert variables['kernel'].sharding == ('din', 'dout')
+        assert variables['bias'].value.shape == (3,)
+        assert variables['bias'].sharding == ('dout',)
 
         return x, None
 
@@ -481,15 +481,15 @@ class TestScan:
     m = MLP(rngs=nnx.Rngs(0))
 
     # test sharding layers axes is set
-    state = m.get_state()
-    assert state.scan_module.linear.variables.kernel.value.shape == (5, 3, 3)
-    assert state.scan_module.linear.variables.kernel.sharding == (
+    variables = m.get_state().variables
+    assert variables['scan_module']['linear']['kernel'].value.shape == (5, 3, 3)
+    assert variables['scan_module']['linear']['kernel'].sharding == (
       'layers',
       'din',
       'dout',
     )
-    assert state.scan_module.linear.variables.bias.value.shape == (5, 3)
-    assert state.scan_module.linear.variables.bias.sharding == (
+    assert variables['scan_module']['linear']['bias'].value.shape == (5, 3)
+    assert variables['scan_module']['linear']['bias'].sharding == (
       'layers',
       'dout',
     )
@@ -498,15 +498,15 @@ class TestScan:
     y, out = m(x, None)
 
     # test sharding axes is preserved
-    state = m.get_state()
-    assert state.scan_module.linear.kernel.shape == (5, 3, 3)
-    assert state.scan_module.linear.variables.kernel.sharding == (
+    variables = m.get_state().variables
+    assert variables['scan_module']['linear']['kernel'].value.shape == (5, 3, 3)
+    assert variables['scan_module']['linear']['kernel'].sharding == (
       'layers',
       'din',
       'dout',
     )
-    assert state.scan_module.linear.bias.shape == (5, 3)
-    assert state.scan_module.linear.variables.bias.sharding == (
+    assert variables['scan_module']['linear']['bias'].value.shape == (5, 3)
+    assert variables['scan_module']['linear']['bias'].sharding == (
       'layers',
       'dout',
     )