[nnx] add compat

flax/experimental/nnx/__init__.py

@@ -18,10 +18,11 @@
 from flax.linen.pooling import pool as pool
 from flax.typing import Initializer as Initializer
 
-from .nnx import compatibility as compatibility
+from .nnx.compat import wrappers as wrappers
 from .nnx import graph as graph
 from .nnx import errors as errors
-from .nnx import errors as helpers
+from .nnx import helpers as helpers
+from .nnx import compat as compat
 from .nnx.filterlib import All as All
 from .nnx.filterlib import Not as Not
 from .nnx.graph import GraphDef as GraphDef

flax/experimental/nnx/nnx/compat/__init__.py

@@ -0,0 +1,9 @@
+
+from .module import ModuleMeta as ModuleMeta
+from .module import Module as Module
+from .module import Scope as Scope
+from .module import compact as compact
+from .wrappers import functional as functional
+from .wrappers import LinenWrapper as LinenWrapper
+from .wrappers import Functional as Functional
+from .wrappers import NNXWrapper as NNXWrapper

flax/experimental/nnx/nnx/compat/module.py

@@ -0,0 +1,226 @@
+# 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.
+
+from __future__ import annotations
+
+from collections import defaultdict
+import dataclasses
+import functools
+import threading
+import typing as tp
+import typing_extensions as tpe
+
+from flax.experimental.nnx.nnx import graph, rnglib
+import flax.experimental.nnx.nnx.module as nnx_module
+from flax.experimental.nnx.nnx.proxy_caller import (
+  CallableProxy,
+  DelayedAccessor,
+)
+from flax.experimental.nnx.nnx.object import Object
+
+M = tp.TypeVar('M', bound='Module')
+F = tp.TypeVar('F', bound=tp.Callable[..., tp.Any])
+
+
+@dataclasses.dataclass
+class CompactContext:
+  module: 'Module'
+  type_counter: defaultdict[type, int] = dataclasses.field(
+    default_factory=lambda: defaultdict(lambda: 0)
+  )
+
+
+@dataclasses.dataclass
+class ModuleContext(threading.local):
+  parent_stack: list[tp.Optional[CompactContext]] = dataclasses.field(
+    default_factory=lambda: [None]
+  )
+
+
+MODULE_CONTEXT = ModuleContext()
+
+
+@dataclasses.dataclass
+class Scope(Object):
+  rngs: rnglib.Rngs
+
+
+@tp.runtime_checkable
+class _HasSetup(tp.Protocol):
+  def setup(self) -> None: ...
+
+
+class ModuleMeta(nnx_module.ModuleMeta):
+  if not tp.TYPE_CHECKING:
+
+    def __call__(cls, *args, **kwargs):
+      return _module_meta_call(cls, *args, **kwargs)
+
+
+def _module_meta_call(cls: tp.Type[M], *args, **kwargs) -> M:
+  # compact behavior
+  parent_ctx = MODULE_CONTEXT.parent_stack[-1]
+  parent = None
+  module: M
+
+  if parent_ctx is not None:
+    if 'parent' in kwargs:
+      parent = kwargs.pop('parent')
+      if parent is not None:
+        raise ValueError(
+          f"'parent' can only be set to None, got {type(parent).__name__}"
+        )
+      name = None
+    else:
+      type_index = parent_ctx.type_counter[cls]
+      parent_ctx.type_counter[cls] += 1
+
+      # define the name
+      if 'name' in kwargs:
+        name = kwargs.pop('name')
+        if not isinstance(name, str):
+          raise ValueError(f"'name' must be a 'str', got {type(name).__name__}")
+      else:
+        name = f'{cls.__name__}_{type_index}'
+
+      parent = parent_ctx.module
+
+      if hasattr(parent, name):
+        module = getattr(parent, name)
+        return module
+  else:
+    name = None
+
+  module = nnx_module.ModuleMeta.__call__(cls, *args, **kwargs)
+  module.scope = None
+
+  if parent is not None:
+    assert name is not None
+    setattr(parent, name, module)
+    # adopt the parent scope
+    module.scope = parent.scope
+
+  if dataclasses.is_dataclass(module):
+    if isinstance(module, _HasSetup):
+      module.setup()
+
+  return module
+
+
+class ModuleBase:
+  if tp.TYPE_CHECKING:
+    scope: Scope | None
+
+
+@tpe.dataclass_transform(field_specifiers=(dataclasses.field,))  # type: ignore[not-supported-yet]
+class Module(nnx_module.Module, ModuleBase, metaclass=ModuleMeta):
+  def _set_scope(self, scope: Scope | None):
+    """Recursively sets the scope for the Module and its children."""
+    for _, value in graph.iter_graph(self):
+      if isinstance(value, Module):
+        value.scope = scope
+
+  @property
+  def init(self: M) -> M:
+    """Calls a method in initialization mode.
+
+    When a method is called using ``init``, the ``is_initializing`` method
+    will return ``True``. This is useful to implement Modules that support
+    lazy initialization.
+
+    Example::
+
+      >>> from flax.experimental import nnx
+      >>> from flax.experimental.nnx import compat as nnc
+      >>> import jax
+      >>> import jax.numpy as jnp
+      ...
+      >>> class Linear(nnc.Module):
+      ...   def __init__(self, dout, rngs: nnx.Rngs):
+      ...     self.dout = dout
+      ...     self.rngs = rngs
+      ...
+      ...   def __call__(self, x):
+      ...     if self.is_initializing():
+      ...       din = x.shape[-1]
+      ...       if not hasattr(self, 'w'):
+      ...         key = self.rngs.params()
+      ...         self.w = nnx.Param(jax.random.uniform(key, (din, self.dout)))
+      ...       if not hasattr(self, 'b'):
+      ...         self.b = nnx.Param(jnp.zeros((self.dout,)))
+      ...
+      ...     return x @ self.w + self.b
+      ...
+      >>> linear = Linear(3, nnx.Rngs(0))
+      >>> x = jnp.ones((5, 2))
+      >>> y = linear.init(x)
+      >>> linear.w.value.shape
+      (2, 3)
+      >>> linear.b.value.shape
+      (3,)
+      >>> y.shape
+      (5, 3)
+    """
+
+    def _init_context(accessor: DelayedAccessor, *args, **kwargs):
+      for _, value in graph.iter_graph(self):
+        if isinstance(value, Object):
+          value._object__state._initializing = True
+
+      method = accessor(self)
+      try:
+        out = method(*args, **kwargs)
+      finally:
+        for _, value in graph.iter_graph(self):
+          if isinstance(value, Object):
+            value._object__state._initializing = False
+
+      return out
+
+    return CallableProxy(_init_context)  # type: ignore
+
+  def is_initializing(self) -> bool:
+    """Returns whether the Module is initializing.
+
+    ``is_initializing`` returns ``True`` if the Module is currently being run
+    under ``init``.
+    """
+
+    return self._object__state._initializing
+
+  def __init_subclass__(cls, experimental_pytree: bool = False) -> None:
+    super().__init_subclass__(experimental_pytree)
+
+    cls = dataclasses.dataclass(repr=False)(cls)
+
+
+def compact(f: F) -> F:
+  @functools.wraps(f)
+  def compact_wrapper(self, *args, **kwargs):
+    if not isinstance(self, Module):
+      raise ValueError(
+        f"Expected 'self' to be a nnx.compat.Module, got {type(self).__name__}"
+      )
+
+    MODULE_CONTEXT.parent_stack.append(CompactContext(self))
+
+    try:
+      return f(self, *args, **kwargs)
+    finally:
+      MODULE_CONTEXT.parent_stack.pop()
+
+  return compact_wrapper  # type: ignore
+
+
+# register Module as a dataclass_transform

flax/experimental/nnx/nnx/module.py

@@ -14,7 +14,6 @@
 
 from __future__ import annotations
 
-import dataclasses
 import typing as tp
 from functools import partial
 
@@ -27,10 +26,6 @@
 from flax.experimental.nnx.nnx import variables as variableslib
 from flax.experimental.nnx.nnx.graph import GraphDef
 from flax.experimental.nnx.nnx.object import Object, ObjectMeta
-from flax.experimental.nnx.nnx.proxy_caller import (
-  CallableProxy,
-  DelayedAccessor,
-)
 from flax.experimental.nnx.nnx.state import State, StateLeaf
 from flax.typing import Path, PathParts
 
@@ -39,34 +34,17 @@
 M = tp.TypeVar('M', bound='Module')
 S = tp.TypeVar('S', bound=tp.Union[State, tuple[State, ...]])
 V = tp.TypeVar('V', bound=variableslib.Variable[tp.Any])
+F = tp.TypeVar('F', bound=tp.Callable[..., tp.Any])
 
 StateMapping = tp.Mapping[Path, tp.Any]
 tuple_reduce = lambda xs, x: xs + (x,)
 tuple_init = lambda: ()
 
 
-@tp.runtime_checkable
-class _HasSetup(tp.Protocol):
-  def setup(self) -> None: ...
-
-
 class ModuleMeta(ObjectMeta):
-  if not tp.TYPE_CHECKING:
-
-    def __call__(cls, *args: Any, **kwargs: Any) -> Any:
-      return _module_meta_call(cls, *args, **kwargs)
-
-
-def _module_meta_call(cls: tp.Type[M], *args, **kwargs) -> M:
-  module: M = ObjectMeta.__call__(cls, *args, **kwargs)
-
-  if dataclasses.is_dataclass(module):
-    if isinstance(module, _HasSetup):
-      module.setup()
-
-  assert isinstance(module, Module)
-
-  return module
+  # we keep a trivial derived class just in case we need to
+  # add more functionality in the future
+  pass
 
 
 class Module(Object, metaclass=ModuleMeta):
@@ -96,73 +74,6 @@ def sow(
       reduced_value = reduce_fn(init_fn(), value)
       setattr(self, name, variable_type(reduced_value))
 
-  @property
-  def init(self: M) -> M:
-    """Calls a method in initialization mode.
-
-    When a method is called using ``init``, the ``is_initializing`` method
-    will return ``True``. This is useful to implement Modules that support
-    lazy initialization.
-
-    Example::
-
-      >>> from flax.experimental import nnx
-      >>> import jax
-      >>> import jax.numpy as jnp
-      ...
-      >>> class Linear(nnx.Module):
-      ...   def __init__(self, dout, rngs: nnx.Rngs):
-      ...     self.dout = dout
-      ...     self.rngs = rngs
-      ...
-      ...   def __call__(self, x):
-      ...     if self.is_initializing():
-      ...       din = x.shape[-1]
-      ...       if not hasattr(self, 'w'):
-      ...         key = self.rngs.params()
-      ...         self.w = nnx.Param(jax.random.uniform(key, (din, self.dout)))
-      ...       if not hasattr(self, 'b'):
-      ...         self.b = nnx.Param(jnp.zeros((self.dout,)))
-      ...
-      ...     return x @ self.w + self.b
-      ...
-      >>> linear = Linear(3, nnx.Rngs(0))
-      >>> x = jnp.ones((5, 2))
-      >>> y = linear.init(x)
-      >>> linear.w.value.shape
-      (2, 3)
-      >>> linear.b.value.shape
-      (3,)
-      >>> y.shape
-      (5, 3)
-    """
-
-    def _init_context(accessor: DelayedAccessor, *args, **kwargs):
-      for _, value in graph.iter_graph(self):
-        if isinstance(value, Object):
-          value._object__state._initializing = True
-
-      method = accessor(self)
-      try:
-        out = method(*args, **kwargs)
-      finally:
-        for _, value in graph.iter_graph(self):
-          if isinstance(value, Object):
-            value._object__state._initializing = False
-
-      return out
-
-    return CallableProxy(_init_context)  # type: ignore
-
-  def is_initializing(self) -> bool:
-    """Returns whether the Module is initializing.
-
-    ``is_initializing`` returns ``True`` if the Module is currently being run
-    under ``init``.
-    """
-
-    return self._object__state._initializing
-
   def iter_modules(self) -> tp.Iterator[tuple[PathParts, Module]]:
     """Iterates over all nested Modules of the current Module, including the current Module.
 

flax/experimental/nnx/nnx/object.py

@@ -70,11 +70,14 @@ class ObjectMeta(ABCMeta):
     def __call__(cls, *args: Any, **kwargs: Any) -> Any:
       return _graph_node_meta_call(cls, *args, **kwargs)
 
+  def _object_meta_construct(cls, self, *args, **kwargs):
+    self.__init__(*args, **kwargs)
+
 
 def _graph_node_meta_call(cls: tp.Type[G], *args, **kwargs) -> G:
   node = cls.__new__(cls, *args, **kwargs)
   vars(node)['_object__state'] = ObjectState()
-  node.__init__(*args, **kwargs)  # type: ignore[misc]
+  cls._object_meta_construct(node, *args, **kwargs)
 
   return node