From 72ea9599cb371f76e5bca23129603ed0be4e674f Mon Sep 17 00:00:00 2001 From: Cristian Garcia Date: Mon, 13 May 2024 14:12:21 +0100 Subject: [PATCH] [nnx] add compat --- flax/experimental/nnx/__init__.py | 5 +- flax/experimental/nnx/nnx/compat/__init__.py | 9 + flax/experimental/nnx/nnx/compat/module.py | 226 ++++++++++++++++++ .../{compatibility.py => compat/wrappers.py} | 0 flax/experimental/nnx/nnx/module.py | 97 +------- flax/experimental/nnx/nnx/object.py | 5 +- .../nnx/tests/compat/test_module.py | 134 +++++++++++ .../test_wrappers.py} | 7 +- flax/experimental/nnx/tests/test_module.py | 42 ---- 9 files changed, 383 insertions(+), 142 deletions(-) create mode 100644 flax/experimental/nnx/nnx/compat/__init__.py create mode 100644 flax/experimental/nnx/nnx/compat/module.py rename flax/experimental/nnx/nnx/{compatibility.py => compat/wrappers.py} (100%) create mode 100644 flax/experimental/nnx/tests/compat/test_module.py rename flax/experimental/nnx/tests/{test_compatibility.py => compat/test_wrappers.py} (85%) diff --git a/flax/experimental/nnx/__init__.py b/flax/experimental/nnx/__init__.py index 7beeea336..a9b2c8062 100644 --- a/flax/experimental/nnx/__init__.py +++ b/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 diff --git a/flax/experimental/nnx/nnx/compat/__init__.py b/flax/experimental/nnx/nnx/compat/__init__.py new file mode 100644 index 000000000..9505b88b6 --- /dev/null +++ b/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 diff --git a/flax/experimental/nnx/nnx/compat/module.py b/flax/experimental/nnx/nnx/compat/module.py new file mode 100644 index 000000000..c152811a1 --- /dev/null +++ b/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 diff --git a/flax/experimental/nnx/nnx/compatibility.py b/flax/experimental/nnx/nnx/compat/wrappers.py similarity index 100% rename from flax/experimental/nnx/nnx/compatibility.py rename to flax/experimental/nnx/nnx/compat/wrappers.py diff --git a/flax/experimental/nnx/nnx/module.py b/flax/experimental/nnx/nnx/module.py index 806794626..1cb578d83 100644 --- a/flax/experimental/nnx/nnx/module.py +++ b/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. diff --git a/flax/experimental/nnx/nnx/object.py b/flax/experimental/nnx/nnx/object.py index e961c8794..1f86f5dca 100644 --- a/flax/experimental/nnx/nnx/object.py +++ b/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 diff --git a/flax/experimental/nnx/tests/compat/test_module.py b/flax/experimental/nnx/tests/compat/test_module.py new file mode 100644 index 000000000..70bd403c5 --- /dev/null +++ b/flax/experimental/nnx/tests/compat/test_module.py @@ -0,0 +1,134 @@ +# 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 dataclasses + +import jax +import jax.numpy as jnp + +from flax.experimental import nnx +from flax.experimental.nnx import compat + + +class TestCompatModule: + def test_compact_basic(self): + class Linear(compat.Module): + dout: int + + def setup(self): + self.count = 0 + + def __call__(self, x): + self.count += 1 + if not hasattr(self, 'w'): + assert self.scope is not None + rngs = self.scope.rngs + self.w = nnx.Param( + jax.random.uniform(rngs(), (x.shape[-1], self.dout)) + ) + self.b = nnx.Param(jnp.zeros((self.dout,))) + return x @ self.w + self.b[None] + + @dataclasses.dataclass + class Foo(compat.Module): + dout: int + + @compat.compact + def __call__(self, x): + din = x.shape[-1] + self.linear = Linear(self.dout) + x = self.linear(x) + return x + + foo = Foo(5) + x = jnp.ones((3, 2)) + rngs = nnx.Rngs(0) + + foo._set_scope(compat.Scope(rngs)) + y = foo(x) + foo._set_scope(None) + + assert y.shape == (3, 5) + assert hasattr(foo, 'Linear_0') + + assert foo.linear is foo.Linear_0 + assert foo.linear.count == 1 + assert rngs.default.count.value == 1 + + foo._set_scope(compat.Scope(rngs)) + y = foo(x) + foo._set_scope(None) + + assert foo.linear is foo.Linear_0 + assert foo.linear.count == 2 + + # Rngs not called again + assert rngs.default.count.value == 1 + + def test_compact_parent_none(self): + class Foo(compat.Module): + pass + + class Bar(compat.Module): + @compat.compact + def __call__(self): + return Foo().scope + + rngs = nnx.Rngs(0) + bar = Bar() + bar._set_scope(compat.Scope(rngs)) + scope = bar() + bar._set_scope(None) + assert bar.scope is None + assert scope.rngs is rngs + + class Baz(compat.Module): + @compat.compact + def __call__(self): + return Foo(parent=None).scope + + baz = Baz() + baz._set_scope(compat.Scope(rngs)) + scope = baz() + baz._set_scope(None) + assert scope is None + + def test_name(self): + class Foo(compat.Module): + dout: int + + def __call__(self, x): + if not hasattr(self, 'w'): + assert self.scope is not None + rngs = self.scope.rngs + self.w = nnx.Param( + jax.random.uniform(rngs(), (x.shape[-1], self.dout)) + ) + return x @ self.w + + class Bar(compat.Module): + @compat.compact + def __call__(self, x): + return Foo(5, name='foo')(x) + + bar = Bar() + x = jnp.ones((1, 2)) + rngs = nnx.Rngs(0) + bar._set_scope(compat.Scope(rngs)) + y = bar(x) + bar._set_scope(None) + assert y.shape == (1, 5) + + assert hasattr(bar, 'foo') + assert isinstance(bar.foo, Foo) \ No newline at end of file diff --git a/flax/experimental/nnx/tests/test_compatibility.py b/flax/experimental/nnx/tests/compat/test_wrappers.py similarity index 85% rename from flax/experimental/nnx/tests/test_compatibility.py rename to flax/experimental/nnx/tests/compat/test_wrappers.py index af20eedd3..1b5cd2bf7 100644 --- a/flax/experimental/nnx/tests/test_compatibility.py +++ b/flax/experimental/nnx/tests/compat/test_wrappers.py @@ -16,12 +16,13 @@ from flax import linen from flax.experimental import nnx +from flax.experimental.nnx import compat class TestCompatibility: def test_functional(self): # Functional API for NNX Modules - functional = nnx.compatibility.functional(nnx.Linear)(32, 64) + functional = compat.functional(nnx.Linear)(32, 64) state = functional.init(rngs=nnx.Rngs(0)) x = jax.numpy.ones((1, 32)) y, updates = functional.apply(state)(x) @@ -30,7 +31,5 @@ def test_linen_wrapper(self): ## Wrapper API for Linen Modules linen_module = linen.Dense(features=64) x = jax.numpy.ones((1, 32)) - module = nnx.compatibility.LinenWrapper( - linen_module, x, rngs=nnx.Rngs(0) - ) # init + module = compat.LinenWrapper(linen_module, x, rngs=nnx.Rngs(0)) # init y = module(x) # apply diff --git a/flax/experimental/nnx/tests/test_module.py b/flax/experimental/nnx/tests/test_module.py index 856fcc34d..2446c3d73 100644 --- a/flax/experimental/nnx/tests/test_module.py +++ b/flax/experimental/nnx/tests/test_module.py @@ -471,31 +471,6 @@ def __init__(self, din, dout, *, rngs: nnx.Rngs): raise_if_not_found=False, ) - def test_init(self): - 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[None] - - linear = Linear(3, nnx.Rngs(0)) - x = jnp.ones((5, 2)) - y = linear.init(x) - assert linear.w.value.shape == (2, 3) - assert linear.b.value.shape == (3,) - assert y.shape == (5, 3) - assert not linear.is_initializing() - - class TestModulePytree: def test_tree_map(self): class Foo(nnx.Module, experimental_pytree=True): @@ -581,23 +556,6 @@ def __call__(self, x): assert hasattr(m, 'bar') - def test_setup_is_called(self): - @dataclasses.dataclass - class DFoo(nnx.Module): - din: int - dout: int - rngs: nnx.Rngs - - def setup(self): - self.bar = nnx.Linear(self.din, self.dout, rngs=self.rngs) - - def __call__(self, x): - return self.bar(x) - - m = DFoo(1, 1, rngs=nnx.Rngs(0)) - - assert hasattr(m, 'bar') - class TestModuleDef: def test_apply(self):