add dataclass

PiperOrigin-RevId: 755993917

Author: Cristian Garcia

flax/nnx/__init__.py

@@ -175,3 +175,6 @@
 from .extract import NodeStates as NodeStates
 from .summary import tabulate as tabulate
 from . import traversals as traversals
+from .dataclasses import dataclass as dataclass
+from .dataclasses import Static as Static
+from .dataclasses import field as field

flax/nnx/dataclasses.py

@@ -0,0 +1,152 @@
+# 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
+
+import dataclasses
+from dataclasses import field
+import typing as tp
+import typing_extensions as tpe
+
+from flax import config
+from flax.nnx import object as objectlib
+
+A = tp.TypeVar('A')
+T = tp.TypeVar('T', bound=type[objectlib.Object])
+
+
+class StaticTag:
+  ...
+
+
+Static = tp.Annotated[A, StaticTag]
+
+
+def _is_static(annotation):
+  return getattr(annotation, '__metadata__', None) == (StaticTag,)
+
+
+@tp.overload
+def dataclass(
+    cls: T,
+    /,
+    *,
+    init: bool = True,
+    eq: bool = True,
+    order: bool = False,
+    unsafe_hash: bool = False,
+    match_args: bool = True,
+    kw_only: bool = False,
+    slots: bool = False,
+) -> T:
+  ...
+
+
+@tp.overload
+def dataclass(
+    *,
+    init: bool = True,
+    eq: bool = True,
+    order: bool = False,
+    unsafe_hash: bool = False,
+    match_args: bool = True,
+    kw_only: bool = False,
+    slots: bool = False,
+) -> tp.Callable[[T], T]:
+  ...
+
+
+@tpe.dataclass_transform(
+    field_specifiers=(field,),
+)
+def dataclass(
+    cls: T | None = None,
+    /,
+    *,
+    init: bool = True,
+    eq: bool = True,
+    order: bool = False,
+    unsafe_hash: bool = False,
+    match_args: bool = True,
+    kw_only: bool = False,
+    slots: bool = False,
+) -> T | tp.Callable[[T], T]:
+  """Makes an nnx.Object type as a dataclass and defines its pytree node attributes using type hints.
+
+  ``nnx.dataclass`` can be used to create pytree dataclass types using type
+  hints instead of the ``__data__`` attribute. By default, all fields are
+  considered to be nodes, to mark a field as static annotate it with
+  ``nnx.Static[T]``.
+
+  Example::
+
+    >>> from flax import nnx
+    >>> import jax
+    ...
+    >>> @nnx.dataclass
+    ... class Foo(nnx.Object):
+    ...   a: int
+    ...   b: jax.Array
+    ...   c: nnx.Static[int]
+    ...
+    >>> tree = Foo(a=1, b=jax.numpy.array(1), c=1)
+    >>> assert len(jax.tree.leaves(tree)) == 2 # a and b
+
+  ``dataclass`` will raise a ``ValueError`` if the class does not derive from
+  ``nnx.Object``, if the parent Object has ``pytree`` set to anything other than
+  ``'strict'``, or if the class has a ``__data__`` attribute.
+
+  ``nnx.dataclass`` doesn't accept ``repr`` and defines it as ``False`` to avoid
+  overwriting the default ``__repr__`` method from ``Object``.
+  """
+
+  def _dataclass(cls: T):
+    if not issubclass(cls, objectlib.Object):
+      raise ValueError(
+          'dataclass can only be used with a class derived from nnx.Object'
+      )
+    if '__data__' in vars(cls):
+      raise ValueError(
+          'dataclass can only be used with a class without a __data__ attribute'
+      )
+    if config.flax_mutable_array:
+      if cls._object__pytree_mode != 'strict':
+        raise ValueError(
+            "dataclass can only be used with a class with pytree='strict', "
+            f'got {cls._object__pytree_mode}'
+        )
+
+      # here we redefine _object__nodes using the type hints
+      hints = cls.__annotations__
+      all_nodes = list(cls._object__nodes)
+      all_nodes.extend(name for name, typ in hints.items() if not _is_static(typ))
+      cls._object__nodes = frozenset(all_nodes)
+
+    cls = dataclasses.dataclass( # type: ignore
+        cls,
+        init=init,
+        repr=False,
+        eq=eq,
+        order=order,
+        unsafe_hash=unsafe_hash,
+        match_args=match_args,
+        kw_only=kw_only,
+        slots=slots,
+    )
+    return cls
+
+  if cls is None:
+    return _dataclass
+  else:
+    return _dataclass(cls)

flax/nnx/object.py

@@ -42,7 +42,6 @@
 O = tp.TypeVar('O', bound='Object')
 
 BUILDING_DOCS = 'FLAX_DOC_BUILD' in os.environ
-DEFAULT_PYTREE_MODE = 'strict' if config.flax_mutable_array else None
 
 def _collect_stats(
   node: tp.Any, node_stats: dict[int, dict[type[Variable], SizeBytes]]
@@ -192,9 +191,8 @@ class Object(reprlib.Representable, metaclass=ObjectMeta):
     _object__state: ObjectState
 
   def __init_subclass__(
-    cls, pytree: tp.Literal['strict', 'auto', 'all'] | None = DEFAULT_PYTREE_MODE, **kwargs
+    cls, pytree: tp.Literal['strict', 'auto', 'all'] = 'strict', **kwargs
   ) -> None:
-    cls._object__pytree_mode = pytree
     super().__init_subclass__(**kwargs)
 
     graph.register_graph_node_type(
@@ -207,6 +205,7 @@ def __init_subclass__(
       init=cls._graph_node_init,  # type: ignore
     )
     if config.flax_mutable_array and pytree is not None:
+      cls._object__pytree_mode = pytree
       parent_pytree_mode = getattr(cls, '_object__pytree_mode', None)
       if (
           parent_pytree_mode is not None
@@ -463,7 +462,7 @@ def _object__strict_unflatten(
     vars_obj.update(zip(node_names, node_attrs, strict=True))
     vars_obj.update(static_attrs)
     return obj
-  
+
   # all
   def _object__all_flatten_with_paths(self):
     obj_vars = vars(self)

tests/nnx/graph_utils_test.py

@@ -27,7 +27,7 @@
 class List(nnx.Module):
   if config.flax_mutable_array:
     __data__ = ('items',)
-  
+
   def __init__(self, items):
     self.items = list(items)
 

tests/nnx/mutable_array_test.py

@@ -12,20 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pytest
-
-import flax.errors
+from absl.testing import absltest
 from flax import config
 from flax import nnx
+import flax.errors
 import jax
 import jax.numpy as jnp
-from absl.testing import absltest
+import pytest
+
 
 @pytest.mark.skipif(
   not config.flax_mutable_array, reason='MutableArray not enabled'
 )
-class TestMutableArray(absltest.TestCase):
-  def test_tree_map(self):
+class TestObject(absltest.TestCase):
+  def test_pytree(self):
     class Foo(nnx.Module):
       __data__ = ('node',)
 
@@ -40,6 +40,25 @@ def __init__(self):
     assert m.node == 2
     assert m.meta == 1
 
+  def test_pytree_dataclass(self):
+    @nnx.dataclass
+    class Foo(nnx.Module):
+      node: jax.Array
+      meta: nnx.Static[int]
+
+    m = Foo(node=jnp.array(1), meta=1)
+
+    m: Foo = jax.tree.map(lambda x: x + 1, m)
+
+    assert m.node == 2
+    assert m.meta == 1
+
+
+@pytest.mark.skipif(
+    not config.flax_mutable_array, reason='MutableArray not enabled'
+)
+class TestMutableArray(absltest.TestCase):
+
   def test_static(self):
     class C(nnx.Module):
       def __init__(self, meta):
@@ -185,3 +204,7 @@ def test_rngs_call(self):
     rngs = nnx.Rngs(0)
     key = rngs()
     self.assertIsInstance(key, jax.Array)
+
+
+if __name__ == '__main__':
+  absltest.main()