remove VariableState

Author: cgarciae

docs_nnx/api_reference/flax.nnx/variables.rst

@@ -18,8 +18,6 @@ variables
    :members:
 .. autoclass:: VariableMetadata
    :members:
-.. autoclass:: VariableState
-   :members:
 
 .. autofunction:: with_metadata
 

docs_nnx/guides/checkpointing.ipynb

@@ -131,7 +131,7 @@
     "\n",
     "## Restore checkpoints\n",
     "\n",
-    "Note that you saved the checkpoint as a Flax class of [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State), which is also nested with [`nnx.VariableState`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.VariableState) and [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) classes.\n",
+    "Note that you saved the checkpoint as a Flax class of [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State), which is also nested with [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable) and [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) classes.\n",
     "\n",
     "At checkpoint restoration time, you need to have these classes ready in your runtime, and instruct the checkpointing library (Orbax) to restore your pytree back to that structure. This can be achieved as follows:\n",
     "- First, create an abstract Flax NNX model (without allocating any memory for arrays), and show its abstract variable state to the checkpointing library.\n",

docs_nnx/guides/checkpointing.md

@@ -82,7 +82,7 @@ checkpointer.save(ckpt_dir / 'state', state)
 
 ## Restore checkpoints
 
-Note that you saved the checkpoint as a Flax class of [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State), which is also nested with [`nnx.VariableState`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.VariableState) and [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) classes.
+Note that you saved the checkpoint as a Flax class of [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State), which is also nested with [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable) and [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) classes.
 
 At checkpoint restoration time, you need to have these classes ready in your runtime, and instruct the checkpointing library (Orbax) to restore your pytree back to that structure. This can be achieved as follows:
 - First, create an abstract Flax NNX model (without allocating any memory for arrays), and show its abstract variable state to the checkpointing library.

docs_nnx/guides/filters_guide.ipynb

@@ -415,7 +415,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

docs_nnx/nnx_glossary.rst

@@ -37,7 +37,4 @@ For additional terms, refer to the `JAX glossary <https://jax.readthedocs.io/en/
       A Flax NNX transformation (transform) is a wrapped version of a `JAX transformation <https://flax.readthedocs.io/en/latest/guides/transforms.html>`__ that allows the function that is being transformed to take the Flax NNX :term:`Module<Module>` (``nnx.Module``) as input or output. For example, a "lifted" version of `jax.jit <https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html#jax.jit>`__ is :meth:`nnx.jit <flax.nnx.jit>`. Check out the `Flax NNX transforms guide <https://flax.readthedocs.io/en/latest/guides/transforms.html>`__ to learn more.
 
     Variable
-      The weights / parameters / data / array :class:`nnx.Variable <flax.nnx.Variable>` residing in a Flax :term:`Module<Module>`. Variables are defined inside modules as :class:`nnx.Variable <flax.nnx.Variable>` or its subclasses.
-
-    Variable state
-      :class:`nnx.VariableState <flax.nnx.VariableState>` is a purely functional `JAX pytree <https://jax.readthedocs.io/en/latest/working-with-pytrees.html>`__ of all the :term:`Variables<Variable>` inside a :term:`Module<Module>`. Since it is pure, it can be an input or output of a `JAX transformation <https://jax.readthedocs.io/en/latest/key-concepts.html#transformations>`__ function. ``nnx.VariableState`` is obtained by using :meth:`nnx.split <flax.nnx.split>` on the :class:`nnx.Module <flax.nnx.Module>`. (Refer to  :term:`splitting<Split and merge>` and :term:`Module<Module>` to learn more.)
+      The weights / parameters / data / array :class:`nnx.Variable <flax.nnx.Variable>` residing in a Flax :term:`Module<Module>`. Variables are defined inside modules as :class:`nnx.Variable <flax.nnx.Variable>` or its subclasses.

examples/nnx_toy_examples/10_fsdp_and_optimizer.py

@@ -71,49 +71,58 @@ def __call__(self, x: jax.Array):
     return nnx.relu(x @ self.w1 + self.b1) @ self.w2
 
 
-class SGDState(nnx.Variable):
+class SGDState(nnx.Variable[jax.Array]):
   pass
 
 
 class SGD(nnx.Object):
-  def __init__(self, params: nnx.State, lr, decay=0.9):
+  def __init__(self, params, lr, decay=0.9):
     def init_optimizer_state(variable: nnx.Variable):
       return SGDState(
         jnp.zeros_like(variable.value), **variable.get_metadata()
       )
 
     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,
+      params,
+      is_leaf=lambda x: isinstance(x, nnx.Variable),
+    )
     self.decay = decay
 
-  def update(self, grads: nnx.State):
+  def update(self, params, grads):
     def update_fn(
-      params: nnx.Variable, momentum: SGDState, grad: nnx.VariableState
+      params: nnx.Variable[jax.Array], momentum: SGDState, grad: jax.Array
     ):
       # v_t = β * v_{t-1} + (1 - β) * ∇J(θ_t)
-      momentum.value = self.decay * momentum + (1 - self.decay) * grad.value
+      momentum.value = 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.value -= self.lr * momentum[...]
+
+    jax.tree.map(
+      update_fn,
+      params,
+      self.momentum,
+      grads,
+      is_leaf=lambda x: isinstance(x, nnx.Variable),
+    )
 
 
 @nnx.jit
 def create_model():
   model = MLP(1, 32, 1, rngs=nnx.Rngs(0))
   optimizer = SGD(nnx.variables(model, nnx.Param), 0.01, decay=0.9)
-  state = nnx.state(optimizer)
+  state = nnx.variables(optimizer)
   sharded_state = jax.lax.with_sharding_constraint(
     state, nnx.get_named_sharding(state, mesh)
   )
 
-  def get_named_shardings(path: tuple, value: nnx.VariableState):
+  def get_named_shardings(path: tuple, value: nnx.Variable):
     if path[0] == 'params':
-      return value.replace(NamedSharding(mesh, P(*value.sharding)))
+      return NamedSharding(mesh, P(*value.sharding))
     elif path[0] == 'momentum':
       # currently the same as above but in general it could be different
-      return value.replace(NamedSharding(mesh, P(*value.sharding)))
+      return NamedSharding(mesh, P(*value.sharding))
     else:
       raise ValueError(f'Unknown path: {path}')
 
@@ -137,7 +146,7 @@ def loss_fn(model):
     return loss
 
   loss, grad = nnx.value_and_grad(loss_fn)(model)
-  optimizer.update(grad)
+  optimizer.update(nnx.variables(model, nnx.Param), grad)
   return loss
 
 

flax/nnx/__init__.py

@@ -43,6 +43,7 @@
 from .graph import split as split
 from .graph import update as update
 from .graph import clone as clone
+from .graph import pure as pure
 from .graph import pop as pop
 from .graph import state as state
 from .graph import graphdef as graphdef
@@ -161,7 +162,6 @@
 from .variablelib import Intermediate as Intermediate
 from .variablelib import Perturbation as Perturbation
 from .variablelib import Variable as Variable
-from .variablelib import VariableState as VariableState
 from .variablelib import VariableMetadata as VariableMetadata
 from .variablelib import with_metadata as with_metadata
 from .variablelib import variable_type_from_name as variable_type_from_name

flax/nnx/bridge/module.py

@@ -374,28 +374,27 @@ def _get_variables(self) -> tp.Mapping:
     state = graph.state(self)
     _variables: dict = {}
 
-    variable_state: variablelib.VariableState
-    for path, variable_state in statelib.to_flat_state(state):
-
-      if issubclass(variable_state.type, rnglib.RngState):
+    variable: variablelib.Variable
+    for path, variable in statelib.to_flat_state(state):
+      if isinstance(variable, rnglib.RngState):
         # Don't return RNG states, since Linen doesn't have them.
         continue
 
       try:
-        collection = variablelib.variable_name_from_type(variable_state.type)
+        collection = variablelib.variable_name_from_type(type(variable))
       except ValueError:
-        collection = variable_state.type.__name__
+        collection = type(variable).__name__
 
       if collection not in _variables:
         _variables[collection] = {}
 
       if (
-        isinstance(variable_state, variablelib.VariableState)
-        and not variable_state._var_metadata
+        isinstance(variable, variablelib.Variable)
+        and not variable._var_metadata
       ):
-        leaf = variable_state.value
+        leaf = variable.value
       else:
-        leaf = bridge_variables.to_linen_var(variable_state)
+        leaf = bridge_variables.to_linen_var(variable)
 
       _variables[collection][path] = leaf
 

flax/nnx/bridge/variables.py

@@ -43,7 +43,7 @@ def _variable_parents_count(t: type):
 
 
 class NNXMeta(struct.PyTreeNode, meta.AxisMetadata[A]):
-  """Default Flax metadata class for `nnx.VariableState`."""
+  """Default Flax metadata class for `nnx.Variable`."""
 
   var_type: type[variablelib.Variable[tp.Any]] = struct.field(pytree_node=False)
   value: Any = struct.field(pytree_node=True)
@@ -65,14 +65,14 @@ def remove_axis(self, index: int, params: dict[Any, Any]) -> 'NNXMeta[A]':
 
   def get_partition_spec(self) -> jax.sharding.PartitionSpec:
     """Returns the ``Partitionspec`` for this partitioned value."""
-    nnx_var = self.to_nnx_variable().to_state()
-    return spmd.get_partition_spec(nnx_var).value
+    nnx_var = self.to_nnx_variable()
+    return spmd.get_partition_spec(nnx_var)
 
   def to_nnx_variable(self) -> variablelib.Variable:
     return self.var_type(self.value, **self.metadata)
 
 
-def is_vanilla_variable(vs: variablelib.VariableState) -> bool:
+def is_vanilla_variable(vs: variablelib.Variable) -> bool:
   """A variables state is vanilla if its metadata is essentially blank.
 
   Returns False only if it has non-empty hooks or any non-built-in attribute.
@@ -86,7 +86,7 @@ def is_vanilla_variable(vs: variablelib.VariableState) -> bool:
   return True
 
 
-def to_linen_var(vs: variablelib.VariableState) -> meta.AxisMetadata:
+def to_linen_var(vs: variablelib.Variable) -> meta.AxisMetadata:
   metadata = vs.get_metadata()
   if 'linen_meta_type' in metadata:
     linen_type = metadata['linen_meta_type']
@@ -95,7 +95,7 @@ def to_linen_var(vs: variablelib.VariableState) -> meta.AxisMetadata:
     return linen_type(vs.value, **metadata)
   if is_vanilla_variable(vs):
     return vs.value
-  return NNXMeta(vs.type, vs.value, metadata)
+  return NNXMeta(type(vs), vs.value, metadata)
 
 
 def get_col_name(keypath: tp.Sequence[Any]) -> str:
@@ -150,10 +150,7 @@ def nnx_attrs_to_linen_vars(nnx_attrs: dict) -> dict:
   for kp, v in traversals.flatten_mapping(nnx_attrs).items():
     if isinstance(v, variablelib.Variable):
       col_name = variablelib.variable_name_from_type(type(v))
-      v = to_linen_var(v.to_state())
-    elif isinstance(v, variablelib.VariableState):
-      col_name = variablelib.variable_name_from_type(v.type)
-      v = to_linen_var(v)
+      v = to_linen_var(v.copy())
     elif isinstance(v, graph.GraphDef):
       col_name = 'nnx'  # an nnx.GraphDef for some ToLinen submodule
     else:

flax/nnx/bridge/wrappers.py

@@ -239,8 +239,8 @@ class ToLinen(linen.Module):
   args: tp.Sequence = ()
   kwargs: tp.Mapping[str, tp.Any] = FrozenDict({})
   skip_rng: bool = False
-  metadata_fn: tp.Callable[[variablelib.VariableState], tp.Any] | None = (
-      bv.to_linen_var
+  metadata_fn: tp.Callable[[variablelib.Variable], tp.Any] | None = (
+    bv.to_linen_var
   )
 
   @linen.compact
@@ -297,7 +297,7 @@ def _update_variables(self, module):
 
     # group state by collection
     for path, leaf in nnx.to_flat_state(state):
-      type_ = leaf.type if isinstance(leaf, nnx.VariableState) else type(leaf)
+      type_ = type(leaf)
       collection = variablelib.variable_name_from_type(
           type_, allow_register=True
       )
@@ -310,7 +310,7 @@ def _update_variables(self, module):
       if self.is_mutable_collection(collection):
 
         def _to_linen_var(x):
-          if isinstance(x, nnx.VariableState):
+          if isinstance(x, nnx.Variable):
             if self.metadata_fn:
               return self.metadata_fn(x)
             else:
@@ -319,22 +319,22 @@ def _to_linen_var(x):
 
         collection_state = nnx.traversals.unflatten_mapping(flat_state)
         collection_state = jax.tree.map(
-            _to_linen_var,
-            collection_state,
-            is_leaf=lambda x: isinstance(x, nnx.VariableState),
+          _to_linen_var,
+          collection_state,
+          is_leaf=lambda x: isinstance(x, nnx.Variable),
         )
         for k, v in collection_state.items():
           self.put_variable(collection, k, v)
 
 
 def to_linen(
-    nnx_class: tp.Callable[..., Module],
-    *args,
-    metadata_fn: (
-        tp.Callable[[variablelib.VariableState], tp.Any] | None
-    ) = bv.to_linen_var,
-    name: str | None = None,
-    **kwargs,
+  nnx_class: tp.Callable[..., Module],
+  *args,
+  metadata_fn: (
+    tp.Callable[[variablelib.Variable], tp.Any] | None
+  ) = bv.to_linen_var,
+  name: str | None = None,
+  **kwargs,
 ):
   """Shortcut of `nnx.bridge.ToLinen` if user is not changing any of its default fields."""
   return ToLinen(

flax/nnx/extract.py

@@ -179,7 +179,7 @@ def from_prefixes(
 def default_split_fn(
   ctx: graph.SplitContext, path: KeyPath, prefix: Prefix, leaf: Leaf
 ) -> tp.Any:
-  return NodeStates.from_split(*ctx.split(leaf))
+  return NodeStates.from_split(*graph.pure(ctx.split(leaf)))
 
 
 def to_tree(