move freeze, mutable, and pure to graph.py

Author: cgarciae

examples/nnx_toy_examples/mutable_array_demo.py

@@ -31,88 +31,6 @@
 # # Utils
 A = TypeVar('A')
 
-def mutable_like(path, x):
-  return (isinstance(x, nnx.Variable) and x.mutable) or nnx.is_mutable_array(x)
-
-
-def freeze(x: A, only: nnx.filterlib.Filter = mutable_like) -> A:
-  freeze_filter = nnx.filterlib.to_predicate(only)
-  mutable_arrays: set[int] = set()
-
-  def check_mutable_array(path, x):
-    m_array_id = id(x)
-    if m_array_id in mutable_arrays:
-      path_str = jax.tree_util.keystr(path)
-      raise ValueError(
-        f'Found duplicate MutableArray found at path {path_str}: {x}'
-      )
-    mutable_arrays.add(m_array_id)
-
-  def _freeze_fn(jax_path, x):
-    path = tuple(nnx.graph._key_path_to_key(part) for part in jax_path)
-    if freeze_filter(path, x):
-      if isinstance(x, nnx.Variable):
-        check_mutable_array(jax_path, x.raw_value)
-        return x.from_metadata(x[...], x.get_metadata().copy())
-      elif nnx.is_mutable_array(x):
-        check_mutable_array(jax_path, x)
-        return x[...]
-    return x
-
-  return jax.tree.map_with_path(
-    _freeze_fn, x, is_leaf=lambda x: isinstance(x, nnx.Variable)
-  )
-
-
-def array_like(path, x):
-  return (
-    isinstance(x, nnx.Variable) and not x.mutable
-  ) or nnx.is_mutable_array(x)
-
-
-def mutable(x: A, only: nnx.filterlib.Filter = array_like) -> A:
-  freeze_filter = nnx.filterlib.to_predicate(only)
-  mutable_arrays: dict[int, Any] = {}
-
-  def get_mutable(x):
-    m_array_id = id(x)
-    if m_array_id in mutable_arrays:
-      return mutable_arrays[m_array_id]
-
-    if isinstance(x, nnx.Variable):
-      assert not x.mutable
-      _mutable = x.from_metadata(
-        nnx.mutable_array(x.raw_value),
-        x.get_metadata().copy(),
-      )
-      mutable_arrays[m_array_id] = _mutable
-      return _mutable
-    elif isinstance(x, jax.Array):
-      _mutable = nnx.mutable_array(x)
-      mutable_arrays[m_array_id] = _mutable
-      return _mutable
-    return x
-
-  def _mutable_fn(jax_path, x):
-    path = tuple(nnx.graph._key_path_to_key(part) for part in jax_path)
-    if freeze_filter(path, x):
-      return get_mutable(x)
-    return x
-
-  return jax.tree.map_with_path(
-    _mutable_fn, x, is_leaf=lambda x: isinstance(x, nnx.Variable)
-  )
-
-def pure(tree: A) -> A:
-  def _pure_fn(x):
-    if isinstance(x, nnx.Variable | nnx.VariableState):
-      return x.raw_value
-    return x
-
-  return jax.tree.map(
-    _pure_fn, tree, is_leaf=lambda x: isinstance(x, nnx.Variable | nnx.VariableState)
-  )
-
 def fork_rngs(
   rngs: nnx.Rngs,
   /,

flax/nnx/__init__.py

@@ -53,6 +53,9 @@
 from .graph import MergeContext as MergeContext
 from .graph import merge_context as merge_context
 from .graph import variables as variables
+from .graph import freeze as freeze
+from .graph import mutable as mutable
+from .graph import pure as pure
 from .graph import cached_partial as cached_partial
 from .nn import initializers as initializers
 from .nn.activations import celu as celu
@@ -168,6 +171,7 @@
 from .variablelib import variable_name_from_type as variable_name_from_type
 from .variablelib import register_variable_name as register_variable_name
 from .variablelib import mutable_array as mutable_array
+from .variablelib import MutableArray as MutableArray
 from .variablelib import is_mutable_array as is_mutable_array
 from .visualization import display as display
 from .extract import to_tree as to_tree