[nnx] add optax optimizer

Author: cgarciae

.github/workflows/flax_test.yml

@@ -79,6 +79,30 @@ jobs:
     - name: Test importing Flax
       run: |
         uv run python -c "import flax"
+  test-mutable-array:
+    name: Run MutableArray tests
+    needs: [pre-commit, commit-count, test-import]
+    runs-on: ubuntu-24.04-16core
+    steps:
+    - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+    - name: Set up Python 3.11
+      id: setup_python
+      uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
+      with:
+        python-version: 3.11
+    - name: Setup uv
+      uses: astral-sh/setup-uv@887a942a15af3a7626099df99e897a18d9e5ab3a # v5.1.0
+      with:
+        version: "0.3.0"
+    - name: Install dependencies
+      run: |
+        uv sync --extra all --extra testing --extra docs
+        uv pip install -U git+https://github.com/jax-ml/jax.git
+    - name: Run MutableArray tests
+      run: |
+        source .venv/bin/activate
+        FLAX_MUTABLE_ARRAY=true pytest tests/nnx/mutable_array_test.py
+
   tests:
     name: Run Tests
     needs: [pre-commit, commit-count, test-import]

flax/nnx/__init__.py

@@ -139,6 +139,7 @@
 from .training.metrics import Metric as Metric
 from .training.metrics import MultiMetric as MultiMetric
 from .training.optimizer import Optimizer as Optimizer
+from .training.optimizer import OptaxOptimizer as OptaxOptimizer
 from .transforms.autodiff import DiffState as DiffState
 from .transforms.autodiff import grad as grad
 from .transforms.autodiff import value_and_grad as value_and_grad

flax/nnx/graph.py

@@ -2652,10 +2652,22 @@ def clone(node: Node) -> Node:
   graphdef, state = split(node)
   return merge(graphdef, state)
 
+def find_duplicates(tree) -> tuple[str, str] | None:
+  mutable_arrays: dict[int, str] = {}
+  paths_leaves = jax.tree.leaves_with_path(tree)
+  for path, x in paths_leaves:
+    m_array_id = id(x)
+    if m_array_id in mutable_arrays:
+      current_path_str = jax.tree_util.keystr(path)
+      previous_path_str = mutable_arrays[m_array_id]
+      return current_path_str, previous_path_str
+    mutable_arrays[m_array_id] = jax.tree_util.keystr(path)
+
+  return None
 
 def _mutable_like(path, x):
   return (
-    isinstance(x, Variable) and x.mutable
+    isinstance(x, Variable | VariableState) and x.mutable
   ) or variablelib.is_mutable_array(x)
 
 
@@ -2681,7 +2693,7 @@ def freeze(tree: A, /, only: filterlib.Filter = _mutable_like) -> A:
     ...   nnx.freeze(tree)
     ... except ValueError as e:
     ...   print(e)
-    Found duplicate MutableArray found at path [1] and [0] ...
+    Found duplicate at path '[1]' and '[0]'.
 
   ``only`` is a `Filter <https://flax.readthedocs.io/en/latest/guides/filters_guide.html>`__
   that can be used to specify which mutable arrays to freeze::
@@ -2698,45 +2710,36 @@ def freeze(tree: A, /, only: filterlib.Filter = _mutable_like) -> A:
   Returns:
     A pytree with the frozen arrays.
   """
+  if (duplicate := find_duplicates(tree)) is not None:
+    current_path_str, previous_path_str = duplicate
+    raise ValueError(
+      f"Found duplicate at path '{current_path_str}' "
+      f"and '{previous_path_str}'."
+    )
   freeze_filter = filterlib.to_predicate(only)
-  mutable_arrays: dict[int, str] = {}
-
-  def check_mutable_array(path, x):
-    m_array_id = id(x)
-    if m_array_id in mutable_arrays:
-      current_path_str = jax.tree_util.keystr(path)
-      previous_path_str = mutable_arrays[m_array_id]
-      raise ValueError(
-        f'Found duplicate MutableArray found at path {current_path_str} '
-        f'and {previous_path_str} at object {x}.'
-      )
-    mutable_arrays[m_array_id] = jax.tree_util.keystr(path)
 
   def _freeze_fn(jax_path, x):
-    path = tuple(_key_path_to_key(part) for part in jax_path)
+    path = jax_to_nnx_path(jax_path)
     if freeze_filter(path, x):
-      if isinstance(x, Variable):
-        check_mutable_array(jax_path, x.raw_value)
-        return x.from_metadata(x[...], x.get_metadata().copy())
-      elif variablelib.is_mutable_array(x):
-        check_mutable_array(jax_path, x)
-        return x[...]
+      x = jax.tree.map(lambda x: x[...], x)
+    elif isinstance(x, Variable | VariableState):
+      x = jax.tree.map(lambda x: x, x)
     return x
 
   tree = jax.tree.map_with_path(
-    _freeze_fn, tree, is_leaf=lambda x: isinstance(x, Variable)
+    _freeze_fn, tree, is_leaf=lambda x: isinstance(x, Variable | VariableState)
   )
   return tree
 
 
 def _array_like(path, x):
   return (
-    isinstance(x, Variable) and isinstance(x.raw_value, jax.Array)
+    isinstance(x, Variable | VariableState) and not x.mutable
   ) or isinstance(x, jax.Array)
 
 
 def mutable(tree: A, /, only: filterlib.Filter = _array_like) -> A:
-  """Converts a pytree of arrays to mutable arrays.
+  """Converts a tree of arrays to mutable arrays.
 
   Example::
 
@@ -2757,7 +2760,7 @@ def mutable(tree: A, /, only: filterlib.Filter = _array_like) -> A:
     ...   nnx.mutable(tree)
     ... except ValueError as e:
     ...   print(e)
-    Found duplicate Array found at path [1] and [0] ...
+    Found duplicate at path '[1]' and '[0]'.
 
   ``only`` is a `Filter <https://flax.readthedocs.io/en/latest/guides/filters_guide.html>`__
   that can be used to specify which arrays to convert to mutable arrays.
@@ -2774,34 +2777,24 @@ def mutable(tree: A, /, only: filterlib.Filter = _array_like) -> A:
   Returns:
     A pytree with the mutable arrays.
   """
+  if (duplicate := find_duplicates(tree)) is not None:
+    current_path_str, previous_path_str = duplicate
+    raise ValueError(
+      f"Found duplicate at path '{current_path_str}' "
+      f"and '{previous_path_str}'."
+    )
   mutable_filter = filterlib.to_predicate(only)
-  arrays: dict[int, str] = {}
-
-  def check_array(path, x):
-    m_array_id = id(x)
-    if m_array_id in arrays:
-      current_path_str = jax.tree_util.keystr(path)
-      previous_path_str = arrays[m_array_id]
-      raise ValueError(
-        f'Found duplicate Array found at path {current_path_str} '
-        f'and {previous_path_str} at object {x}.'
-      )
-    arrays[m_array_id] = jax.tree_util.keystr(path)
 
   def _mutable_fn(jax_path, x):
-    path = tuple(_key_path_to_key(part) for part in jax_path)
+    path = jax_to_nnx_path(jax_path)
     if mutable_filter(path, x):
-      if isinstance(x, Variable) and isinstance(x.raw_value, jax.Array):
-        check_array(jax_path, x.raw_value)
-        mutable_array = variablelib.mutable_array(x.raw_value)
-        return x.from_metadata(mutable_array, x.get_metadata().copy())
-      elif isinstance(x, jax.Array):
-        check_array(jax_path, x)
-        return variablelib.mutable_array(x)
+      x = jax.tree.map(variablelib.mutable_array, x)
+    elif isinstance(x, Variable | VariableState):
+      x = jax.tree.map(lambda x: x, x)
     return x
 
   return jax.tree.map_with_path(
-    _mutable_fn, tree, is_leaf=lambda x: isinstance(x, Variable)
+    _mutable_fn, tree, is_leaf=lambda x: isinstance(x, Variable | VariableState)
   )
 
 
@@ -3047,6 +3040,11 @@ def _key_path_to_key(key: tp.Any) -> Key:
   else:
     return str(key)
 
+
+def jax_to_nnx_path(jax_path: tuple, /):
+  return tuple(_key_path_to_key(part) for part in jax_path)
+
+
 class IndexesPytreeDef(tp.NamedTuple):
   key_index: HashableMapping[Key, int]
   treedef: jax.tree_util.PyTreeDef

flax/nnx/nn/normalization.py

@@ -18,7 +18,7 @@
 import jax.numpy as jnp
 from jax import lax
 
-from flax import nnx
+from flax import nnx, config
 from flax.nnx import rnglib
 from flax.nnx.module import Module, first_from
 from flax.nnx.nn import dtypes, initializers
@@ -360,11 +360,18 @@ def __call__(
         use_fast_variance=self.use_fast_variance,
         mask=mask,
       )
+      # stop_gradient only for flax_mutable_array
+      if config.flax_mutable_array:
+        stop_gradient = jax.lax.stop_gradient
+      else:
+        stop_gradient = lambda x: x
 
-      self.mean[...] = (
+      self.mean[...] = stop_gradient(
         self.momentum * self.mean.value + (1 - self.momentum) * mean
       )
-      self.var[...] = self.momentum * self.var.value + (1 - self.momentum) * var
+      self.var[...] = stop_gradient(
+        self.momentum * self.var.value + (1 - self.momentum) * var
+      )
 
     return _normalize(
       x,

flax/nnx/training/optimizer.py

@@ -20,8 +20,7 @@
 import optax
 
 from flax import nnx
-from flax.nnx import filterlib
-from flax.nnx import variablelib
+from flax.nnx import filterlib, graph
 from flax.nnx.object import Object
 from flax.nnx.variablelib import Variable, VariableState
 
@@ -51,7 +50,7 @@ class OptVariable(OptState):
 
 def _wrap_optimizer_state(opt_state):
   def wrap_optimizer_state_fn(x):
-    if isinstance(x, variablelib.VariableState):
+    if isinstance(x, VariableState):
       new_state = x.copy()
       new_state.source_type = x.type
       new_state.type = OptVariable
@@ -62,7 +61,7 @@ def wrap_optimizer_state_fn(x):
   return jax.tree.map(
     wrap_optimizer_state_fn,
     opt_state,
-    is_leaf=lambda x: isinstance(x, variablelib.VariableState),
+    is_leaf=lambda x: isinstance(x, VariableState),
   )
 
 
@@ -274,4 +273,98 @@ def update(self, grads, **kwargs):
 
     self.step.value += 1
     nnx.update(self.model, new_params)
-    _update_opt_state(self.opt_state, new_opt_state)
+    _update_opt_state(self.opt_state, new_opt_state)
+
+
+def to_opt_state(tree):
+  def _to_opt_state(x):
+    if isinstance(x, Variable | VariableState):
+      opt_state = OptVariable(x[...], **x.get_metadata())  # type: ignore
+    else:
+      opt_state = OptArray(x)
+    return opt_state
+
+  tree = jax.tree.map(
+    _to_opt_state,
+    tree,
+    is_leaf=lambda x: isinstance(x, Variable | VariableState),
+  )
+  return tree
+
+
+class OptaxOptimizer(Object):
+  """Stateful wrapper around an Optax optimizer.
+
+  Example usage::
+
+    >>> from flax import config
+    >>> if not config.flax_mutable_array:
+    ...   import pytest
+    ...   pytest.skip('MutableArrays required for this example')
+    ...
+    >>> import jax, jax.numpy as jnp
+    >>> from flax import nnx
+    >>> from flax import config
+    >>> import optax
+    ...
+    >>> class Model(nnx.Module):
+    ...   __data__ = ('linear1', 'linear2', 'bn')
+    ...   def __init__(self, rngs):
+    ...     self.linear1 = nnx.Linear(2, 3, rngs=rngs)
+    ...     self.bn = nnx.BatchNorm(3, rngs=rngs)
+    ...     self.linear2 = nnx.Linear(3, 4, rngs=rngs)
+    ...   def __call__(self, x):
+    ...     return self.linear2(nnx.relu(self.bn(self.linear1(x))))
+    ...
+    >>> x = jax.random.normal(jax.random.key(0), (5, 2))
+    >>> y = jnp.ones((5, 4))
+    ...
+    >>> model = Model(nnx.Rngs(1))
+    >>> optimizer = nnx.OptaxOptimizer(nnx.state(model, nnx.Param), tx=optax.adam(1e-3))
+    ...
+    >>> @jax.jit
+    ... def train_step(model, optimizer, x, y):
+    ...   graphdef, params, nondiff = nnx.split(model, nnx.Param, ...)
+    ...   def loss_fn(params):
+    ...     model = nnx.merge(graphdef, params, nondiff)
+    ...     return ((model(x) - y) ** 2).mean()
+    ...
+    ...   loss, grads = jax.value_and_grad(loss_fn)(nnx.freeze(params))
+    ...   optimizer.update(params, grads)
+    ...   return loss
+    ...
+    >>> loss = train_step(model, optimizer, x, y)
+    >>> loss
+    Array(1.2029127, dtype=float32)
+
+  Args:
+    params: The parameters to be optimized.
+    tx: An optax gradient transformation.
+  """
+  __nodes__ = ('step', 'opt_state')
+
+  def __init__(self, params, tx: optax.GradientTransformation):
+    self.tx = tx
+    self.step = OptArray(jnp.array(0, dtype=jnp.uint32))
+    self.opt_state = to_opt_state(tx.init(params))
+
+  def update(self, params, grads, **kwargs):
+    param_arrays = graph.freeze(graph.pure(params))
+    grad_arrays = graph.freeze(graph.pure(grads))
+    opt_state_arrays = graph.freeze(graph.pure(self.opt_state))
+
+    updates, new_opt_state = self.tx.update(
+      grad_arrays, opt_state_arrays, param_arrays, **kwargs
+    )
+    new_params = optax.apply_updates(param_arrays, updates)
+
+    def _update_variable(param, value):
+      param[...] = value
+
+    jax.tree.map(
+      _update_variable,
+      (params, self.opt_state),
+      (new_params, new_opt_state),
+      is_leaf=lambda x: isinstance(x, Variable | VariableState),
+    )
+    self.step[...] += 1

flax/nnx/variablelib.py

@@ -224,13 +224,16 @@ def state(cls, value: A, **metadata) -> VariableState[A]:
     return cls(value, **metadata).to_state()
 
   @property
-  def mutable(self) -> bool | None:
+  def mutable(self) -> bool:
     if is_mutable_array(self.raw_value):
       return True
     elif isinstance(self.raw_value, jax.Array):
       return False
     else:
-      return None
+      raise ValueError(
+        f'mutable is only supported for jax.Array and MutableArray, '
+        f'got {type(self.raw_value).__name__}'
+      )
 
   def get_metadata(self):
     return self._var_metadata
@@ -972,6 +975,18 @@ def raw_value(self) -> A:
   def raw_value(self, value: A) -> None:
     object.__setattr__(self, 'value', value)
 
+  @property
+  def mutable(self) -> bool:
+    if is_mutable_array(self.raw_value):
+      return True
+    elif isinstance(self.raw_value, jax.Array):
+      return False
+    else:
+      raise ValueError(
+        f'mutable is only supported for jax.Array and MutableArray, '
+        f'got {type(self.raw_value).__name__}'
+      )
+
   def __getattribute__(self, name: str) -> None:
     if name == 'value':
       value = object.__getattribute__(self, 'value')