Repository: google/flax

Issue:
nnx.bridge.ToNNX does not handle the addition of top-level variables
Repro:

```python
import flax.linen as nn
import flax.nnx as nnx


class LinenModule(nn.Module):
    @nn.compact
    def __call__(self):
        if not self.is_initializing() and self.is_mutable_collection("cache"):
            self.put_variable("cache", "x", 0)
        return self.get_variable("cache", "x")


class NNXModule(nnx.Module):
    def __init__(self):
        self.module = nnx.bridge.ToNNX(LinenModule()).lazy_init()

    def __call__(self):
        result1 = self.module(mutable=["cache"])
        assert result1 == 0
        result2 = self.module()
        assert result2 == 0, result2  # fails: result2 is None


module = NNXModule()
module()
```

I believe this is because `self.linen_attributes` is never updated after initialization, see [here](https://github.com/google/flax/blob/1961c1259e6c2d22a0c27de9599476aaa073a553/flax/nnx/bridge/wrappers.py#L157).

Discussion:

Patch:
diff --git a/flax/nnx/bridge/wrappers.py b/flax/nnx/bridge/wrappers.py
index eed4ba2f..ab673644 100644
--- a/flax/nnx/bridge/wrappers.py
+++ b/flax/nnx/bridge/wrappers.py
@@ -148,7 +148,7 @@ class ToNNX(Module):
       out, variables = self.module.init_with_output(_rngs, *args, method=method, **kwargs)
 
       nnx_attrs = bv.linen_vars_to_nnx_attrs(variables)
-      linen_attributes = set()
+      linen_attributes = set(self.linen_attributes)
       for attr_name, value in nnx_attrs.items():
         setattr(self, attr_name, value)
         linen_attributes.add(attr_name)
@@ -167,13 +167,17 @@ class ToNNX(Module):
     if kwargs.get('mutable', False) != False:
       out, updates = out
       nnx_attrs = bv.linen_vars_to_nnx_attrs(updates)
+      linen_attributes = set(self.linen_attributes)
       for attr_name, value in nnx_attrs.items():
+        linen_attributes.add(attr_name)
         if hasattr(self, attr_name) and isinstance(value, dict):
           original_tree = getattr(self, attr_name)
           setattr(self, attr_name, original_tree | value)
         else:
           setattr(self, attr_name, value)
 
+      self.linen_attributes = tuple(linen_attributes)  # make it hashable
+
     return out
 
 
diff --git a/flax/nnx/rnglib.py b/flax/nnx/rnglib.py
index ab9817ac..a100740a 100644
--- a/flax/nnx/rnglib.py
+++ b/flax/nnx/rnglib.py
@@ -80,7 +80,7 @@ RngDict = tp.Union[
 ]
 
 
-class Rngs(Object, tp.Mapping[str, tp.Callable[[], jax.Array]]):
+class Rngs(Object, tp.Mapping[str, RngStream]):
   """NNX rng container class. To instantiate the ``Rngs``, pass
   in an integer, specifying the starting seed. ``Rngs`` can have
   different "streams", allowing the user to generate different
diff --git a/flax/nnx/training/metrics.py b/flax/nnx/training/metrics.py
index 18044902..4facf427 100644
--- a/flax/nnx/training/metrics.py
+++ b/flax/nnx/training/metrics.py
@@ -373,7 +373,7 @@ class MultiMetric(Metric):
     for metric_name in self._metric_names:
       getattr(self, metric_name).update(**updates)
 
-  def compute(self) -> dict[str, tp.Any]:
+  def compute(self) -> dict[str, Metric]:
     """Compute and return the value of all underlying ``Metric``'s. This method
     will return a dictionary, mapping strings (defined by the key-word arguments
     ``**metrics`` passed to the constructor) to the corresponding metric value.


Test patch:
diff --git a/tests/nnx/bridge/wrappers_test.py b/tests/nnx/bridge/wrappers_test.py
index 5b65603a..8c584912 100644
--- a/tests/nnx/bridge/wrappers_test.py
+++ b/tests/nnx/bridge/wrappers_test.py
@@ -217,6 +217,28 @@ class TestCompatibility(absltest.TestCase):
     self.assertEqual(jax.tree.structure(from_top_weights),
                      jax.tree.structure(from_middle_weights))
 
+  def test_adding_new_attributes(self):
+    class LinenModule(nn.Module):
+      @nn.compact
+      def __call__(self):
+        if not self.is_initializing() and self.is_mutable_collection('cache'):
+          self.put_variable('cache', 'x', 0)
+        res = self.get_variable('cache', 'x')
+        return res
+
+    class NNXModule(nnx.Module):
+      def __init__(self):
+        self.module = nnx.bridge.ToNNX(LinenModule()).lazy_init()
+
+      def __call__(self):
+        result1 = self.module(mutable=['cache'])
+        assert result1 == 0
+        result2 = self.module()
+        assert result2 == 0, result2  # fails: result2 is None
+
+    module = NNXModule()
+    module()
+
   ##################
   ### NNXToLinen ###
   ##################