Merge pull request #5066 from google:dataclass

PiperOrigin-RevId: 831131118

Author: Flax Authors

docs_nnx/guides/pytree.ipynb

@@ -395,8 +395,8 @@
    "id": "8055b72c",
    "metadata": {},
    "source": [
-    "### Class Annotations\n",
-    "Annotations can also be added at the type level via `nnx.Static` and `nnx.Data`. This is useful for creating `dataclasses` but the mechanism also works for regular classes."
+    "### Dataclasses\n",
+    "`nnx.Pytree` dataclasses can be created by using the `nnx.dataclass` decorator. To control the status of each field, `nnx.static` and `nnx.data` can be used as `field` specifiers."
    ]
   },
   {
@@ -420,25 +420,58 @@
    "source": [
     "import dataclasses\n",
     "\n",
-    "@dataclasses.dataclass\n",
+    "@nnx.dataclass\n",
     "class Foo(nnx.Pytree):\n",
-    "  i: nnx.Data[int]\n",
-    "  s: nnx.Static[str]\n",
+    "  i: int = nnx.data()\n",
     "  x: jax.Array\n",
     "  a: int\n",
+    "  s: str = nnx.static(default='hi', kw_only=True)\n",
     "\n",
-    "@dataclasses.dataclass\n",
+    "@nnx.dataclass\n",
     "class Bar(nnx.Pytree):\n",
-    "  ls: nnx.Data[list[Foo]]\n",
+    "  ls: list[Foo] = nnx.data()\n",
     "  shapes: list[int]\n",
     "\n",
     "pytree = Bar(\n",
-    "  ls=[Foo(i, \"Hi\" + \"!\" * i, jnp.array(42 * i), hash(i)) for i in range(2)],\n",
+    "  ls=[Foo(i, jnp.array(42 * i), hash(i)) for i in range(2)],\n",
     "  shapes=[8, 16, 32]\n",
     ")\n",
     "pytree_structure(pytree)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "fca51f65",
+   "metadata": {},
+   "source": [
+    "`dataclasses.dataclass` can also be used directly, however type checkers will not handle `nnx.static` and `nnx.data` correctly. To solve this `dataclasses.field` can be used by setting `metadata` with the appropriate entry for `static`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "ff54e732",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "dataclass pytree structure:\n",
+      " - pytree.a = 10\n"
+     ]
+    }
+   ],
+   "source": [
+    "@dataclasses.dataclass\n",
+    "class Bar(nnx.Pytree):\n",
+    "  a: int = dataclasses.field(metadata={'static': False}) # data\n",
+    "  b: str = dataclasses.field(metadata={'static': True})  # static\n",
+    "\n",
+    "pytree = Bar(a=10, b=\"hello\")\n",
+    "pytree_structure(pytree, title='dataclass pytree structure')"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "d6036a0e",
@@ -457,7 +490,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 11,
    "id": "509a517e",
    "metadata": {},
    "outputs": [
@@ -501,7 +534,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 12,
    "id": "98e04ff9",
    "metadata": {},
    "outputs": [
@@ -534,7 +567,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 13,
    "id": "c864d5b1",
    "metadata": {},
    "outputs": [
@@ -568,7 +601,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 14,
    "id": "628698dd",
    "metadata": {},
    "outputs": [
@@ -616,20 +649,20 @@
    "id": "37ee2429",
    "metadata": {},
    "source": [
-    "Checking for `nnx.data` or `nnx.static` annotations stored in inside nested structures that are not `nnx.Pytree` instances:"
+    "Checking for `nnx.data` or `nnx.static` annotations stored inside nested structures that are not `nnx.Pytree` instances:"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 15,
    "id": "f9d69634",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "ValueError: Found unexpected tags {'data', 'static'} on attribute 'Foo.a'. Values from nnx.data(...) and\n",
+      "ValueError: Found unexpected tags {'static', 'data'} on attribute 'Foo.a'. Values from nnx.data(...) and\n",
       "nnx.static(...) should be assigned to nnx.Pytree attributes directly, they should not be inside other structures. Got value of type '<class 'list'>' on Pytree of type '<class '__main__.Foo'>'.\n"
      ]
     }
@@ -656,7 +689,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 16,
    "id": "668db479",
    "metadata": {},
    "outputs": [
@@ -704,7 +737,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 17,
    "id": "32c46ce8",
    "metadata": {},
    "outputs": [
@@ -741,7 +774,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 18,
    "id": "c33e4862",
    "metadata": {},
    "outputs": [
@@ -778,7 +811,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 19,
    "id": "dda51b67",
    "metadata": {},
    "outputs": [
@@ -825,7 +858,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 20,
    "id": "caa01e3b",
    "metadata": {},
    "outputs": [
@@ -864,7 +897,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 21,
    "id": "d2e03753",
    "metadata": {},
    "outputs": [
@@ -929,7 +962,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 22,
    "id": "ca9f58a2",
    "metadata": {},
    "outputs": [
@@ -1003,7 +1036,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 38,
+   "execution_count": 23,
    "id": "41398e14",
    "metadata": {},
    "outputs": [
@@ -1050,20 +1083,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 24,
    "id": "d10effba",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "step = 0, loss = Array(0.7326511, dtype=float32), perturbations = \u001b[38;2;79;201;177mState\u001b[0m\u001b[38;2;255;213;3m({\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
+      "step = 0, loss = Array(0.7326511, dtype=float32), iterm_grads = \u001b[38;2;79;201;177mState\u001b[0m\u001b[38;2;255;213;3m({\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
       "  \u001b[38;2;156;220;254m'xgrad'\u001b[0m\u001b[38;2;212;212;212m: \u001b[0m\u001b[38;2;79;201;177mPerturbation\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 3 (12 B)\u001b[0m\n",
       "    \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mArray([[-0.430146  , -0.14356601,  0.2935633 ]], dtype=float32)\n",
       "  \u001b[38;2;255;213;3m)\u001b[0m\n",
       "\u001b[38;2;255;213;3m})\u001b[0m\n",
-      "step = 1, loss = Array(0.65039134, dtype=float32), perturbations = \u001b[38;2;79;201;177mState\u001b[0m\u001b[38;2;255;213;3m({\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
+      "step = 1, loss = Array(0.65039134, dtype=float32), iterm_grads = \u001b[38;2;79;201;177mState\u001b[0m\u001b[38;2;255;213;3m({\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n",
       "  \u001b[38;2;156;220;254m'xgrad'\u001b[0m\u001b[38;2;212;212;212m: \u001b[0m\u001b[38;2;79;201;177mPerturbation\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 3 (12 B)\u001b[0m\n",
       "    \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mArray([[-0.38535568, -0.11745065,  0.24441527]], dtype=float32)\n",
       "  \u001b[38;2;255;213;3m)\u001b[0m\n",
@@ -1108,7 +1141,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": 25,
    "id": "a9cab639",
    "metadata": {},
    "outputs": [

docs_nnx/guides/pytree.md

@@ -195,31 +195,43 @@ pytree = Bar(1.0, True)
 pytree_structure(pytree)
 ```
 
-### Class Annotations
-Annotations can also be added at the type level via `nnx.Static` and `nnx.Data`. This is useful for creating `dataclasses` but the mechanism also works for regular classes.
+### Dataclasses
+`nnx.Pytree` dataclasses can be created by using the `nnx.dataclass` decorator. To control the status of each field, `nnx.static` and `nnx.data` can be used as `field` specifiers.
 
 ```{code-cell} ipython3
 import dataclasses
 
-@dataclasses.dataclass
+@nnx.dataclass
 class Foo(nnx.Pytree):
-  i: nnx.Data[int]
-  s: nnx.Static[str]
+  i: int = nnx.data()
   x: jax.Array
   a: int
+  s: str = nnx.static(default='hi', kw_only=True)
 
-@dataclasses.dataclass
+@nnx.dataclass
 class Bar(nnx.Pytree):
-  ls: nnx.Data[list[Foo]]
+  ls: list[Foo] = nnx.data()
   shapes: list[int]
 
 pytree = Bar(
-  ls=[Foo(i, "Hi" + "!" * i, jnp.array(42 * i), hash(i)) for i in range(2)],
+  ls=[Foo(i, jnp.array(42 * i), hash(i)) for i in range(2)],
   shapes=[8, 16, 32]
 )
 pytree_structure(pytree)
 ```
 
+`dataclasses.dataclass` can also be used directly, however type checkers will not handle `nnx.static` and `nnx.data` correctly. To solve this `dataclasses.field` can be used by setting `metadata` with the appropriate entry for `static`.
+
+```{code-cell} ipython3
+@dataclasses.dataclass
+class Bar(nnx.Pytree):
+  a: int = dataclasses.field(metadata={'static': False}) # data
+  b: str = dataclasses.field(metadata={'static': True})  # static
+
+pytree = Bar(a=10, b="hello")
+pytree_structure(pytree, title='dataclass pytree structure')
+```
+
 ### Attribute Updates
 
 +++
@@ -281,7 +293,7 @@ except ValueError as e:
   print("ValueError:", e)
 ```
 
-Checking for `nnx.data` or `nnx.static` annotations stored in inside nested structures that are not `nnx.Pytree` instances:
+Checking for `nnx.data` or `nnx.static` annotations stored inside nested structures that are not `nnx.Pytree` instances:
 
 ```{code-cell} ipython3
 class Foo(nnx.Pytree):

flax/nnx/__init__.py

@@ -36,6 +36,7 @@
 from .pytreelib import Object as Object
 from .pytreelib import Data as Data
 from .pytreelib import Static as Static
+from .pytreelib import dataclass as dataclass
 from .pytreelib import data as data
 from .pytreelib import static as static
 from .pytreelib import register_data_type as register_data_type
@@ -209,12 +210,12 @@
 from .summary import tabulate as tabulate
 from . import traversals as traversals
 
-# alias VariableState
-VariableState = Variable
 
 import typing as _tp
 
-if not _tp.TYPE_CHECKING:
+if _tp.TYPE_CHECKING:
+  VariableState = Variable
+else:
   def __getattr__(name):
     if name == "VariableState":
       import warnings
@@ -224,4 +225,5 @@ def __getattr__(name):
           DeprecationWarning,
           stacklevel=2,
       )
+      return Variable
     raise AttributeError(f"Module {__name__} has no attribute '{name}'")