Merge pull request #4627 from google:remove-deprecated

PiperOrigin-RevId: 746110808

Author: Flax Authors

docs_nnx/nnx_basics.ipynb

@@ -375,9 +375,9 @@
     "In the code below notice the following:\n",
     "\n",
     "1. The custom `create_model` function takes in a key and returns an `MLP` object, since you create five keys and use [`nnx.vmap`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.vmap) over `create_model` a stack of 5 `MLP` objects is created.\n",
-    "2. The [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.Scan) is used to iteratively apply each `MLP` in the stack to the input `x`.\n",
-    "3. The [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.Scan) (consciously) deviates from [`jax.lax.scan`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html#jax.lax.scan) and instead mimics [`nnx.vmap`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.vmap), which is more expressive. [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.Scan) allows specifying multiple inputs, the scan axes of each input/output, and the position of the carry.\n",
-    "4. [`State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) updates for the [`nnx.BatchNorm`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/normalization.html#flax.nnx.BatchNorm) and [`nnx.Dropout`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/stochastic.html#flax.nnx.Dropout) layers are automatically propagated by [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.Scan)."
+    "2. The [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.scan) is used to iteratively apply each `MLP` in the stack to the input `x`.\n",
+    "3. The [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.scan) (consciously) deviates from [`jax.lax.scan`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html#jax.lax.scan) and instead mimics [`nnx.vmap`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.vmap), which is more expressive. [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.scan) allows specifying multiple inputs, the scan axes of each input/output, and the position of the carry.\n",
+    "4. [`State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) updates for the [`nnx.BatchNorm`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/normalization.html#flax.nnx.BatchNorm) and [`nnx.Dropout`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/stochastic.html#flax.nnx.Dropout) layers are automatically propagated by [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.scan)."
    ]
   },
   {

docs_nnx/nnx_basics.md

@@ -193,9 +193,9 @@ The next example uses Flax [`nnx.vmap`](https://flax.readthedocs.io/en/latest/ap
 In the code below notice the following:
 
 1. The custom `create_model` function takes in a key and returns an `MLP` object, since you create five keys and use [`nnx.vmap`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.vmap) over `create_model` a stack of 5 `MLP` objects is created.
-2. The [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.Scan) is used to iteratively apply each `MLP` in the stack to the input `x`.
-3. The [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.Scan) (consciously) deviates from [`jax.lax.scan`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html#jax.lax.scan) and instead mimics [`nnx.vmap`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.vmap), which is more expressive. [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.Scan) allows specifying multiple inputs, the scan axes of each input/output, and the position of the carry.
-4. [`State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) updates for the [`nnx.BatchNorm`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/normalization.html#flax.nnx.BatchNorm) and [`nnx.Dropout`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/stochastic.html#flax.nnx.Dropout) layers are automatically propagated by [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.Scan).
+2. The [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.scan) is used to iteratively apply each `MLP` in the stack to the input `x`.
+3. The [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.scan) (consciously) deviates from [`jax.lax.scan`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html#jax.lax.scan) and instead mimics [`nnx.vmap`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.vmap), which is more expressive. [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.scan) allows specifying multiple inputs, the scan axes of each input/output, and the position of the carry.
+4. [`State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) updates for the [`nnx.BatchNorm`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/normalization.html#flax.nnx.BatchNorm) and [`nnx.Dropout`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/stochastic.html#flax.nnx.Dropout) layers are automatically propagated by [`nnx.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.scan).
 
 ```{code-cell} ipython3
 @nnx.vmap(in_axes=0, out_axes=0)

flax/nnx/__init__.py

@@ -136,11 +136,6 @@
 from .training.metrics import Metric as Metric
 from .training.metrics import MultiMetric as MultiMetric
 from .training.optimizer import Optimizer as Optimizer
-from .transforms.deprecated import Jit as Jit
-from .transforms.deprecated import Remat as Remat
-from .transforms.deprecated import Scan as Scan
-from .transforms.deprecated import Vmap as Vmap
-from .transforms.deprecated import Pmap as Pmap
 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/extract.py

@@ -16,7 +16,6 @@
 import typing as tp
 
 import jax
-# from jax._src.tree_util import broadcast_prefix
 
 from flax import struct
 from flax.nnx.object import Object
@@ -32,91 +31,6 @@
 Leaf = tp.Any
 
 
-class ExtractionIndex(struct.PyTreeNode):
-  """Index of a graph node in a Pytree structure."""
-
-  index: Index = struct.field(pytree_node=False)
-
-
-@tp.overload
-def extract_graph_nodes(
-  pytree: A,
-  /,
-  *,
-  validate_fn: tp.Callable[[KeyPath, Prefix, Leaf], None] | None = None,
-) -> tuple[A, tuple[tp.Any, ...]]: ...
-@tp.overload
-def extract_graph_nodes(
-  pytree: A,
-  /,
-  *,
-  prefix: tp.Any,
-  validate_fn: tp.Callable[[KeyPath, Prefix, Leaf], None] | None = None,
-) -> tuple[A, tuple[tp.Any, ...], tuple[tp.Any, ...]]: ...
-def extract_graph_nodes(
-  pytree: A,
-  /,
-  *,
-  prefix: tp.Any = Missing,
-  validate_fn: tp.Callable[[KeyPath, Prefix, Leaf], None] | None = None,
-) -> (
-  tuple[A, tuple[tp.Any, ...]]
-  | tuple[A, tuple[tp.Any, ...], tuple[tp.Any, ...]]
-):
-  """Extracts all graph nodes from a pytree."""
-  nodes: dict[tp.Any, Index] = {}
-  node_prefixes = []
-  leaves = []
-
-  prefix_leaves = broadcast_prefix(
-    prefix,
-    pytree,
-    prefix_is_leaf=lambda x: x is None,
-  )
-  key_leaves, treedef = jax.tree_util.tree_flatten_with_path(pytree)
-
-  assert len(key_leaves) == len(prefix_leaves)
-
-  for (keypath, leaf), prefix_leaf in zip(key_leaves, prefix_leaves):
-    if validate_fn:
-      validate_fn(keypath, prefix_leaf, leaf)
-    if graph.is_graph_node(leaf):
-      if leaf not in nodes:
-        index = nodes[leaf] = len(nodes)
-        node_prefixes.append(prefix_leaf)
-      else:
-        index = nodes[leaf]
-        # check consistent aliasing
-        if prefix_leaf != node_prefixes[index]:
-          path_str = jax.tree_util.keystr(keypath)
-          raise ValueError(
-            f'Inconsistent aliasing detected. Node {type(leaf)} at path {path_str} '
-            f'has different prefixes: {prefix_leaf} and {node_prefixes[index]}.'
-          )
-      leaves.append(ExtractionIndex(index))
-    else:
-      leaves.append(leaf)
-
-  pytree_out = jax.tree.unflatten(treedef, leaves)
-
-  if prefix is Missing:
-    return pytree_out, tuple(nodes)  # type: ignore[bad-return-type]
-  else:
-    return pytree_out, tuple(nodes), tuple(node_prefixes)  # type: ignore[bad-return-type]
-
-
-def insert_graph_nodes(pytree: A, nodes: tuple[tp.Any, ...], /) -> A:
-  """Inserts graph nodes into a pytree."""
-
-  def _maybe_insert(x):
-    if isinstance(x, ExtractionIndex):
-      return nodes[x.index]
-    return x
-
-  return jax.tree.map(
-    _maybe_insert, pytree, is_leaf=lambda x: isinstance(x, ExtractionIndex)
-  )
-
 class PrefixMapping(abc.ABC):
   @abc.abstractmethod
   def map_prefix(
@@ -342,7 +256,7 @@ def from_tree(
 ) -> tp.Any:
   if prefix is Missing or prefix is None:
     # fast path, no need for prefix broadcasting or consistent aliasing checks
-    with graph.merge_context(is_inner, ctxtag) as merge_ctx:
+    with graph.merge_context(ctxtag, is_inner) as merge_ctx:
 
       def maybe_split(x):
         if (
@@ -366,7 +280,7 @@ def maybe_split(x):
   assert len(leaf_keys) == len(leaf_prefixes)
   leaves_out = []
 
-  with graph.merge_context(is_inner, ctxtag) as merge_ctx:
+  with graph.merge_context(ctxtag, is_inner) as merge_ctx:
     for (keypath, leaf), leaf_prefix in zip(leaf_keys, leaf_prefixes):
       if (
         map_non_graph_nodes

flax/nnx/graph.py

@@ -1633,9 +1633,9 @@ def unflatten(
 def merge_context(): ...
 @tp.overload
 @contextlib.contextmanager
-def merge_context(inner: bool | None, ctxtag: tp.Hashable | None): ...
+def merge_context(ctxtag: tp.Hashable | None, inner: bool | None): ...
 @contextlib.contextmanager
-def merge_context(inner: bool | None = None, ctxtag: tp.Hashable | None = None):
+def merge_context(ctxtag: tp.Hashable | None = None, inner: bool | None = None):
   GRAPH_CONTEXT.index_ref_stack.append(MergeContext(ctxtag, {}, inner))
 
   try:
@@ -1691,138 +1691,6 @@ def unflatten_end(self, index_ref: dict[Index, tp.Any], inner_merge: bool):
           (obj, index) for index, obj in index_ref.items()
       )
 
-  @tp.overload
-  def split(self, graph_node: A, /) -> tuple[GraphDef[A], GraphState]:
-    ...
-
-  @tp.overload
-  def split(
-      self, graph_node: A, first: filterlib.Filter, /
-  ) -> tuple[GraphDef[A], GraphState]:
-    ...
-
-  @tp.overload
-  def split(
-      self,
-      graph_node: A,
-      first: filterlib.Filter,
-      second: filterlib.Filter,
-      /,
-      *filters: filterlib.Filter,
-  ) -> tuple[GraphDef[A], GraphState, tpe.Unpack[tuple[GraphState, ...]]]:
-    ...
-
-  def split(
-    self, node: A, *filters: filterlib.Filter
-  ) -> tuple[GraphDef[A], GraphState, tpe.Unpack[tuple[GraphState, ...]]]:
-    """Split a graph node into a :class:`GraphDef` and one or more :class:`State`s. State is
-    a ``Mapping`` from strings or integers to ``Variables``, Arrays or nested States. GraphDef
-    contains all the static information needed to reconstruct a ``Module`` graph, it is analogous
-    to JAX’s ``PyTreeDef``. :func:`split` is used in conjunction with :func:`merge` to switch
-    seamlessly between stateful and stateless representations of the graph.
-
-    Example usage::
-
-      >>> from flax import nnx
-      >>> import jax, jax.numpy as jnp
-      ...
-      >>> class Foo(nnx.Module):
-      ...   def __init__(self, rngs):
-      ...     self.batch_norm = nnx.BatchNorm(2, rngs=rngs)
-      ...     self.linear = nnx.Linear(2, 3, rngs=rngs)
-      ...
-      >>> node = Foo(nnx.Rngs(0))
-      >>> graphdef, params, batch_stats = nnx.split(node, nnx.Param, nnx.BatchStat)
-      ...
-      >>> jax.tree.map(jnp.shape, params)
-      State({
-        'batch_norm': {
-          'bias': VariableState(
-            type=Param,
-            value=(2,)
-          ),
-          'scale': VariableState(
-            type=Param,
-            value=(2,)
-          )
-        },
-        'linear': {
-          'bias': VariableState(
-            type=Param,
-            value=(3,)
-          ),
-          'kernel': VariableState(
-            type=Param,
-            value=(2, 3)
-          )
-        }
-      })
-      >>> jax.tree.map(jnp.shape, batch_stats)
-      State({
-        'batch_norm': {
-          'mean': VariableState(
-            type=BatchStat,
-            value=(2,)
-          ),
-          'var': VariableState(
-            type=BatchStat,
-            value=(2,)
-          )
-        }
-      })
-
-    Arguments:
-      node: graph node to split.
-      *filters: some optional filters to group the state into mutually exclusive substates.
-    Returns:
-      :class:`GraphDef` and one or more :class:`State`'s equal to the number of filters passed. If no
-      filters are passed, a single :class:`State` is returned.
-    """
-    ref_index: RefMap = RefMap()
-    graphdef, flat_state = flatten(
-      node, ref_index=ref_index, ref_outer_index=self.inner_ref_outer_index
-    )
-    flat_states = _split_state(flat_state, filters)
-    states = _to_nested_state(graphdef, flat_states)
-    assert len(states) >= 1
-    self.flatten_end(ref_index)
-    return graphdef, *states  # type: ignore[return-value]
-
-  def merge(
-    self,
-    graphdef: GraphDef[A],
-    state: GraphState,
-    *states: GraphState,
-  ) -> A:
-    """merge"""
-    if not isinstance(graphdef, NodeDef):
-      raise ValueError(
-        f'Expected a NodeDef instance, but got {type(graphdef)}.'
-      )
-    if self.outer_ref_outer_index is None:
-      raise ValueError('Cannot merge without ref_index.')
-
-    if self.outer_ref_outer_index is not None:
-      # outer merge (4), create index_ref_cache
-      index_ref_cache = self.outer_index_outer_ref
-      assert index_ref_cache is not None
-    else:
-      # inner merge (2)
-      index_ref_cache = None
-
-    _state = _merge_to_flat_state((state, *states))
-    index_ref: dict[Index, tp.Any] = {}
-    node = unflatten(
-      graphdef,
-      _state,
-      index_ref=index_ref,
-      outer_index_outer_ref=index_ref_cache,
-    )
-
-    self.unflatten_end(index_ref, True)
-
-    return node
-
 
 jax.tree_util.register_static(UpdateContext)
 
@@ -1919,16 +1787,20 @@ def update_context(tag: tp.Hashable):
     >>> from flax import nnx
     ...
     >>> m1 = nnx.Dict({})
-    >>> with nnx.update_context('example') as ctx:
-    ...   graphdef, state = ctx.split(m1)
+    >>> with nnx.update_context('example'):
+    ...   with nnx.split_context('example') as ctx:
+    ...     graphdef, state = ctx.split(m1)
     ...   @jax.jit
     ...   def f(graphdef, state):
-    ...     m2 = ctx.merge(graphdef, state)
+    ...     with nnx.merge_context('example', inner=True) as ctx:
+    ...       m2 = ctx.merge(graphdef, state)
     ...     m2.a = 1
     ...     m2.ref = m2  # create a reference cycle
-    ...     return ctx.split(m2)
+    ...     with nnx.split_context('example') as ctx:
+    ...       return ctx.split(m2)
     ...   graphdef_out, state_out = f(graphdef, state)
-    ...   m3 = ctx.merge(graphdef_out, state_out)
+    ...   with nnx.merge_context('example', inner=False) as ctx:
+    ...     m3 = ctx.merge(graphdef_out, state_out)
     ...
     >>> assert m1 is m3
     >>> assert m1.a == 1
@@ -1937,36 +1809,36 @@ def update_context(tag: tp.Hashable):
   Note that ``update_context`` takes in a ``tag`` argument which is used
   primarily as a safety mechanism reduce the risk of accidentally using the
   wrong UpdateContext when using :func:`current_update_context` to access the
-  current active context. current_update_context can be used as a way of
-  accessing the current active context without having to pass it as a capture::
+  current active context. ``update_context`` can also be used as a
+  decorator that creates/activates an UpdateContext context for the
+  duration of the function::
 
     >>> from flax import nnx
     ...
     >>> m1 = nnx.Dict({})
     >>> @jax.jit
     ... def f(graphdef, state):
-    ...   ctx = nnx.current_update_context('example')
-    ...   m2 = ctx.merge(graphdef, state)
+    ...   with nnx.merge_context('example', inner=True) as ctx:
+    ...     m2 = ctx.merge(graphdef, state)
     ...   m2.a = 1     # insert static attribute
     ...   m2.ref = m2  # create a reference cycle
-    ...   return ctx.split(m2)
+    ...   with nnx.split_context('example') as ctx:
+    ...     return ctx.split(m2)
     ...
     >>> @nnx.update_context('example')
     ... def g(m1):
-    ...   ctx = nnx.current_update_context('example')
-    ...   graphdef, state = ctx.split(m1)
+    ...   with nnx.split_context('example') as ctx:
+    ...     graphdef, state = ctx.split(m1)
     ...   graphdef_out, state_out = f(graphdef, state)
-    ...   return ctx.merge(graphdef_out, state_out)
+    ...   with nnx.merge_context('example', inner=False) as ctx:
+    ...     return ctx.merge(graphdef_out, state_out)
     ...
     >>> m3 = g(m1)
     >>> assert m1 is m3
     >>> assert m1.a == 1
     >>> assert m1.ref is m1
 
-  As shown in the code above, ``update_context`` can also be used as a
-  decorator that creates/activates an UpdateContext context for the
-  duration of the function. The context can be accessed using
-  :func:`current_update_context`.
+  The context can be accessed using :func:`current_update_context`.
 
   Args:
     tag: A string tag to identify the context.

flax/nnx/transforms/__init__.py

@@ -11,5 +11,3 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-from . import deprecated