fix RNN

Implements Mapping for `StateAxes` and uses `StateAxes` in place of `dict` to fix RNN, this avoids some JAX pytree errors when scanning attributes for data in `nnx.Pytree`.

PiperOrigin-RevId: 800653246

Author: Cristian Garcia

flax/nnx/nn/recurrent.py

@@ -685,7 +685,7 @@ def __init__(
         'Expected rngs to be a jax.Array, int, Rngs, or bool. '
         f'Got {type(rngs)}.'
       )
-    self.state_axes = state_axes or {...: iteration.Carry}  # type: ignore
+    self.state_axes = state_axes or nnx.StateAxes({...: iteration.Carry})  # type: ignore
     self.broadcast_rngs = broadcast_rngs
 
   def __call__(

flax/nnx/transforms/iteration.py

@@ -56,7 +56,7 @@ class Carry:
 # -------------------------------
 
 
-class StateAxes(extract.PrefixMapping):
+class StateAxes(extract.PrefixMapping, tp.Mapping):
 
   def __init__(
     self,
@@ -101,6 +101,15 @@ def __repr__(self):
   def items(self):
     return zip(self.filters, self.axes)
 
+  def __getitem__(self, key):
+    return self.axes[self.filters.index(key)]
+
+  def __iter__(self):
+    return iter(self.filters)
+
+  def __len__(self):
+    return len(self.filters)
+
   def __eq__(self, other):
     return (
         isinstance(other, StateAxes)