[nnx] don't share Rngs

Author: cgarciae

docs_nnx/api_reference/flax.nnx/helpers.rst

@@ -4,10 +4,7 @@ helpers
 .. automodule:: flax.nnx
 .. currentmodule:: flax.nnx
 
-.. autoclass:: Dict
-   :members:
-.. autoclass:: List
-   :members:
+
 .. autoclass:: Sequential
    :members:
 .. autoclass:: TrainState

docs_nnx/api_reference/flax.nnx/index.rst

@@ -7,6 +7,7 @@ Experimental API. See the `NNX page <https://flax.readthedocs.io/en/latest/nnx/i
   :maxdepth: 3
 
   graph
+  object
   module
   nn/index
   rnglib

docs_nnx/api_reference/flax.nnx/object.rst

@@ -0,0 +1,13 @@
+object
+------------------------
+
+.. automodule:: flax.nnx
+.. currentmodule:: flax.nnx
+
+.. autoclass:: Object
+   :members:
+.. autofunction:: data
+.. autodata:: Data
+   :annotation:
+.. autofunction:: is_data_type
+.. autofunction:: register_data_type

docs_nnx/api_reference/flax.nnx/training/optimizer.rst

@@ -6,3 +6,6 @@ Optimizer
 
 .. autoclass:: Optimizer
    :members: __init__, update
+
+.. autoclass:: PytreeOptimizer
+   :members: __init__, update

docs_nnx/api_reference/flax.nnx/variables.rst

@@ -23,4 +23,4 @@ variables
 
 .. autofunction:: variable_name_from_type
 .. autofunction:: variable_type_from_name
-.. autofunction:: set_variable_name
+.. autofunction:: register_variable_name

docs_nnx/guides/filters_guide.ipynb

@@ -5,7 +5,7 @@
    "id": "95b08e64",
    "metadata": {},
    "source": [
-    "# Using Filters, grouping NNX variables \n",
+    "# Filters\n",
     "\n",
     "Flax NNX uses [`Filter`s](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/filterlib.html) extensively as a way to create [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) groups in APIs, such as [`nnx.split`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.split), [`nnx.state()`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.state), and many of the [Flax NNX transformations (transforms)](https://flax.readthedocs.io/en/latest/guides/jax_and_nnx_transforms.html).\n",
     "\n",

docs_nnx/guides/filters_guide.md

@@ -8,7 +8,7 @@ jupytext:
     jupytext_version: 1.13.8
 ---
 
-# Using Filters, grouping NNX variables 
+# Filters
 
 Flax NNX uses [`Filter`s](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/filterlib.html) extensively as a way to create [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) groups in APIs, such as [`nnx.split`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.split), [`nnx.state()`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.state), and many of the [Flax NNX transformations (transforms)](https://flax.readthedocs.io/en/latest/guides/jax_and_nnx_transforms.html).
 

examples/nnx_toy_examples/06_scan_over_layers.py

@@ -47,7 +47,6 @@ def create_block(rngs: nnx.Rngs):
     self.layers = create_block(rngs)
 
   def __call__(self, x: jax.Array) -> jax.Array:
-    @nnx.split_rngs(splits=self.n_layers)
     @nnx.scan
     def scan_fn(x: jax.Array, block: Block):
       x = block(x)

examples/nnx_toy_examples/10_fsdp_and_optimizer.py

@@ -84,19 +84,29 @@ def init_optimizer_state(variable: nnx.Variable):
 
     self.lr = lr
     self.params = params
-    self.momentum: nnx.State = jax.tree.map(init_optimizer_state, self.params)
+    self.momentum: nnx.State = jax.tree.map(
+      init_optimizer_state,
+      self.params,
+      is_leaf=lambda x: isinstance(x, nnx.Variable | nnx.VariableState),
+    )
     self.decay = decay
 
   def update(self, grads: nnx.State):
     def update_fn(
       params: nnx.Variable, momentum: SGDState, grad: nnx.VariableState
     ):
       # v_t = β * v_{t-1} + (1 - β) * ∇J(θ_t)
-      momentum.value = self.decay * momentum + (1 - self.decay) * grad.value
+      momentum[...] = self.decay * momentum[...] + (1 - self.decay) * grad[...]
       # θ_{t+1} = θ_t - α * v_t
-      params.value -= self.lr * momentum
-
-    jax.tree.map(update_fn, self.params, self.momentum, grads)
+      params[...] -= self.lr * momentum[...]
+
+    jax.tree.map(
+      update_fn,
+      self.params,
+      self.momentum,
+      grads,
+      is_leaf=lambda x: isinstance(x, nnx.Variable | nnx.VariableState),
+    )
 
 
 @nnx.jit

examples/nnx_toy_examples/mutable_array_basic.py

@@ -35,8 +35,6 @@ def dataset(batch_size):
 
 
 class Linear(nnx.Module):
-  __data__ = ('w', 'b')
-
   def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):
     self.w = nnx.Param(jax.random.uniform(rngs.params(), (din, dout)))
     self.b = nnx.Param(jnp.zeros((dout,)))
@@ -50,8 +48,6 @@ class Count(nnx.Variable[nnx.A]):
 
 
 class MLP(nnx.Module):
-  __data__ = ('count', 'linear1', 'linear2')
-
   def __init__(self, din, dhidden, dout, *, rngs: nnx.Rngs):
     self.count = Count(jnp.array(0))
     self.linear1 = Linear(din, dhidden, rngs=rngs)