Add deprecations for old NNX view functions

Author: samanklesaria

flax/nnx/__init__.py

@@ -50,7 +50,7 @@
 from .module import M as M
 from .module import Module as Module
 from .module import capture as capture
-from .module import with_modules as with_modules
+from .module import with_modules as with_modules, view as view
 from .module import view_info as view_info
 from .module import with_attributes as with_attributes
 from .module import iter_children as iter_children, iter_modules as iter_modules
@@ -75,8 +75,8 @@
 from .graphlib import MergeContext as MergeContext
 from .graphlib import merge_context as merge_context
 from .graphlib import variables as variables
-from .graphlib import with_vars as with_vars
-from .graphlib import as_pure as as_pure
+from .graphlib import with_vars as with_vars, vars_as as vars_as
+from .graphlib import as_pure as as_pure, pure as pure
 from .graphlib import cached_partial as cached_partial
 from .graphlib import flatten as flatten
 from .graphlib import unflatten as unflatten
@@ -152,7 +152,7 @@
 from .spmd import get_named_sharding as get_named_sharding
 from .spmd import with_partitioning as with_partitioning
 from .spmd import get_abstract_model as get_abstract_model
-from .spmd import as_abstract as as_abstract
+from .spmd import as_abstract as as_abstract, abstract_with_sharding as abstract_with_sharding
 from .statelib import FlatState as FlatState
 from .statelib import State as State
 from .statelib import to_flat_state as to_flat_state

flax/nnx/deprecations.py

@@ -0,0 +1,69 @@
+# Copyright 2024 The Flax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import functools
+import warnings
+from typing import TypeVar
+from collections.abc import Callable
+
+F = TypeVar('F', bound=Callable)
+
+
+def deprecated(new_fn: F) -> F:
+  """Creates a deprecated alias for a renamed function.
+
+  .. deprecated::
+     This decorator is for marking functions as deprecated. The returned
+     wrapper emits a :class:`DeprecationWarning` on every call and then
+     delegates to ``new_fn``.
+
+  The returned callable copies the signature, type annotations, and
+  docstring of ``new_fn``, with a deprecation notice prepended to the
+  docstring. This keeps IDE autocomplete and type-checking working while
+  clearly communicating that callers should migrate.
+
+  Args:
+    new_fn: The current, non-deprecated function to delegate to.
+
+  Returns:
+    A wrapper that emits a :class:`DeprecationWarning` and then calls
+    ``new_fn`` with the same arguments.
+
+  Example::
+
+    >>> from flax.nnx.deprecations import deprecated
+    >>> def new_api(x):
+    ...   return x * 2
+    >>> old_api = deprecated(new_api)
+    >>> old_api(3)  # emits DeprecationWarning: use new_api instead
+    6
+
+  """
+
+  @functools.wraps(new_fn)
+  def wrapper(*args, **kwargs):
+    warnings.warn(
+      f'This function is deprecated. Use {new_fn.__qualname__} instead.',
+      DeprecationWarning,
+      stacklevel=2,
+    )
+    return new_fn(*args, **kwargs)
+
+  dep_notice = (
+    f'.. deprecated::\n'
+    f'   Use :func:`{new_fn.__qualname__}` instead.\n\n'
+  )
+  wrapper.__doc__ = dep_notice + (new_fn.__doc__ or '')
+
+  return wrapper  # type: ignore[return-value]

flax/nnx/graphlib.py

@@ -26,6 +26,7 @@
 from flax import config
 from flax.nnx import filterlib, reprlib, traversals, variablelib
 from flax.nnx import statelib
+from flax.nnx.deprecations import deprecated
 from flax.nnx.proxy_caller import (
   ApplyCaller,
   CallableProxy,
@@ -2761,6 +2762,7 @@ def _to_refs(path, x):
   node = graphlib.map(_to_refs, node, graph=graph)
   return node
 
+vars_as = deprecated(with_vars)
 
 def as_pure(tree: A) -> A:
   """Returns a new tree with all ``Variable`` objects replaced with inner values.
@@ -2808,6 +2810,7 @@ def _pure_fn(_, x):
 
   return map(_pure_fn, tree, auto_create_variables=False)
 
+pure = deprecated(with_modules)
 
 def call(
   graphdef_state: tuple[GraphDef[A], GraphState], /

flax/nnx/module.py

@@ -29,6 +29,7 @@
 from flax.nnx.pytreelib import Pytree, PytreeMeta
 from flax.nnx.graphlib import GraphState
 from flax.nnx.statelib import split_state, State
+from flax.nnx.deprecations import deprecated
 import functools as ft
 from flax.typing import Key, Path, PathParts
 from collections.abc import MutableMapping
@@ -515,6 +516,8 @@ def _set_mode_fn(path, node):
 
   return out
 
+view = deprecated(with_modules)
+
 def with_attributes(
   node: A,
   /,

flax/nnx/spmd.py

@@ -16,6 +16,7 @@
 
 import flax.core.spmd as core_spmd
 from flax.nnx import variablelib, graphlib
+from flax.nnx.deprecations import deprecated
 from flax.nnx.transforms.transforms import eval_shape
 from flax.typing import (
   Sharding,
@@ -238,3 +239,5 @@ def add_sharding(_path, x):
       return abs_var
     return x
   return graphlib.map(add_sharding, tree, graph=graph)
+
+abstract_with_sharding = deprecated(as_abstract)