[nnx] remove VariableState

Author: cgarciae

examples/nnx_toy_examples/10_fsdp_and_optimizer.py

@@ -87,13 +87,13 @@ def init_optimizer_state(variable: nnx.Variable):
     self.momentum: nnx.State = jax.tree.map(
       init_optimizer_state,
       self.params,
-      is_leaf=lambda x: isinstance(x, nnx.Variable | nnx.VariableState),
+      is_leaf=lambda x: isinstance(x, nnx.Variable),
     )
     self.decay = decay
 
   def update(self, grads: nnx.State):
     def update_fn(
-      params: nnx.Variable, momentum: SGDState, grad: nnx.VariableState
+      params: nnx.Variable, momentum: SGDState, grad: nnx.Variable
     ):
       # v_t = β * v_{t-1} + (1 - β) * ∇J(θ_t)
       momentum[...] = self.decay * momentum[...] + (1 - self.decay) * grad[...]
@@ -105,7 +105,7 @@ def update_fn(
       self.params,
       self.momentum,
       grads,
-      is_leaf=lambda x: isinstance(x, nnx.Variable | nnx.VariableState),
+      is_leaf=lambda x: isinstance(x, nnx.Variable),
     )
 
 
@@ -118,12 +118,12 @@ def create_model():
     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}')
 

examples/nnx_toy_examples/mutable_array_demo.py

@@ -192,7 +192,7 @@ def __init__(self, params, lr: float, decay: float = 0.9):
     self.decay = decay
 
     def make_opt_state(x):
-      if isinstance(x, nnx.Variable | nnx.VariableState):
+      if isinstance(x, nnx.Variable):
         return OptState(jnp.zeros_like(x.value), **x.get_metadata())
       else:
         return OptState(jnp.zeros_like(x))
@@ -201,7 +201,7 @@ def make_opt_state(x):
       jax.tree.map(
         make_opt_state,
         params,
-        is_leaf=lambda x: isinstance(x, nnx.Variable | nnx.VariableState),
+        is_leaf=lambda x: isinstance(x, nnx.Variable),
       )
     )
 

flax/nnx/__init__.py

@@ -172,7 +172,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,27 +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,15 +65,17 @@ 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).raw_value
+    nnx_var = self.to_nnx_variable()
+    spec = spmd.get_partition_spec(nnx_var)
+    assert isinstance(spec, jax.sharding.PartitionSpec)
+    return spec
 
   def to_nnx_variable(self) -> variablelib.Variable:
     return self.var_type(self.value, **self.metadata)
 
 
-def is_vanilla_variable(vs: variablelib.VariableState) -> bool:
-  """A variables state is vanilla if its metadata is essentially blank.
+def is_vanilla_variable(vs: variablelib.Variable) -> bool:
+  """A variable 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 +88,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']
@@ -145,14 +147,11 @@ def linen_vars_to_nnx_attrs(variables: tp.Mapping[str, Any]) -> dict[str, Any]:
 
 
 def nnx_attrs_to_linen_vars(nnx_attrs: dict) -> dict:
-  """Convert a dict of NNX variables (or variable states) to Linen-style variables."""
+  """Convert a dict of NNX variables to Linen-style variables."""
   linen_structured = {}
   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)
     elif isinstance(v, graph.GraphDef):
       col_name = 'nnx'  # an nnx.GraphDef for some ToLinen submodule

flax/nnx/bridge/wrappers.py

@@ -252,7 +252,7 @@ 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 = (
+  metadata_fn: tp.Callable[[variablelib.Variable], tp.Any] | None = (
       bv.to_linen_var
   )
 
@@ -310,7 +310,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_ = leaf.type if isinstance(leaf, nnx.Variable) else type(leaf)
       collection = variablelib.variable_name_from_type(
           type_, allow_register=True
       )
@@ -323,7 +323,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:
@@ -334,7 +334,7 @@ def _to_linen_var(x):
         collection_state = jax.tree.map(
             _to_linen_var,
             collection_state,
-            is_leaf=lambda x: isinstance(x, nnx.VariableState),
+            is_leaf=lambda x: isinstance(x, nnx.Variable),
         )
         for k, v in collection_state.items():
           self.put_variable(collection, k, v)
@@ -344,7 +344,7 @@ def to_linen(
     nnx_class: tp.Callable[..., Module],
     *args,
     metadata_fn: (
-        tp.Callable[[variablelib.VariableState], tp.Any] | None
+        tp.Callable[[variablelib.Variable], tp.Any] | None
     ) = bv.to_linen_var,
     name: str | None = None,
     **kwargs,

flax/nnx/filterlib.py

@@ -121,9 +121,7 @@ class OfType:
   type: type
 
   def __call__(self, path: PathParts, x: tp.Any):
-    return isinstance(x, self.type) or (
-      hasattr(x, 'type') and issubclass(x.type, self.type)
-    )
+    return isinstance(x, self.type)
 
   def __repr__(self):
     return f'OfType({self.type!r})'