Repository: jax-ml/jax
Pull request #24700 (2024-11-04 21:24:18)

register_dataclass: allow marking static fields via dataclass.field
This is a stripped-down version of #24664. This PR does not include a jax-specific `dataclass` decorator, but makes `jax.tree_util.register_dataclass` recognize dataclass field metadata for static arguments:
```python
import jax
from dataclasses import dataclass, field

@jax.tree_util.register_dataclass
@dataclass
class MyClass:
    arrs: list[jax.Array]                         # By default, attributes will be treated as data fields.
    name: str = field(metadata={'static': True})  # Optionally mark attributes as static via
                                                  #  `dataclasses.field` mechanism.

m = MyClass([jax.numpy.arange(4), jax.numpy.ones(2)], "my class")

print(jax.tree.structure(m))
# PyTreeDef(CustomNode(MyClass[('my class',)], [[*, *]]))

print(jax.tree.leaves(m))
# [Array([0, 1, 2, 3], dtype=int32), Array([1., 1.], dtype=float32)]


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index bafe55110..0af10bd4c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -43,6 +43,9 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
   * {func}`jax.jit` got a new `compiler_options: dict[str, Any]` argument, for
     passing compilation options to XLA. For the moment it's undocumented and
     may be in flux.
+  * {func}`jax.tree_util.register_dataclass` now allows metadata fields to be
+    declared inline via {func}`dataclasses.field`. See the function documentation
+    for examples.
 
 ## jax 0.4.35 (Oct 22, 2024)
 
diff --git a/jax/_src/deprecations.py b/jax/_src/deprecations.py
index c7a956068..962244a32 100644
--- a/jax/_src/deprecations.py
+++ b/jax/_src/deprecations.py
@@ -125,6 +125,8 @@ register('jax-aval-named-shape')
 register('jax-dlpack-import-legacy')
 register("jax-numpy-astype-complex-to-real")
 register("jax-numpy-array-none")
+register('jax-scipy-beta-args')
+register('tracer-hash')
 register('jax-numpy-reshape-newshape')
 register('jax-numpy-clip-args')
 register('jax-numpy-linalg-matrix_rank-tol')
diff --git a/jax/_src/numpy/reductions.py b/jax/_src/numpy/reductions.py
index be1e55675..fa8d73361 100644
--- a/jax/_src/numpy/reductions.py
+++ b/jax/_src/numpy/reductions.py
@@ -2360,8 +2360,7 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
     index[axis] = high
     high_value = a[tuple(index)]
   else:
-    with jax.debug_nans(False):
-      a = _where(any(lax_internal._isnan(a), axis=axis, keepdims=True), np.nan, a)
+    a = _where(any(lax_internal._isnan(a), axis=axis, keepdims=True), np.nan, a)
     a = lax.sort(a, dimension=axis)
     n = lax.convert_element_type(a_shape[axis], lax_internal._dtype(q))
     q = lax.mul(q, n - 1)
diff --git a/jax/_src/tree_util.py b/jax/_src/tree_util.py
index 474bdfe4e..58f9d7862 100644
--- a/jax/_src/tree_util.py
+++ b/jax/_src/tree_util.py
@@ -927,8 +927,8 @@ def register_pytree_with_keys_class(cls: Typ) -> Typ:
 @export
 def register_dataclass(
     nodetype: Typ,
-    data_fields: Sequence[str],
-    meta_fields: Sequence[str],
+    data_fields: Sequence[str] | None = None,
+    meta_fields: Sequence[str] | None = None,
     drop_fields: Sequence[str] = (),
 ) -> Typ:
   """Extends the set of types that are considered internal nodes in pytrees.
@@ -945,24 +945,33 @@ def register_dataclass(
       attributes represent the whole of the object state, and can be passed
       as keywords to the class constructor to create a copy of the object.
       All defined attributes should be listed among ``meta_fields`` or ``data_fields``.
-    meta_fields: auxiliary data field names. These fields *must* contain static,
-      hashable, immutable objects, as these objects are used to generate JIT cache
-      keys. In particular, ``meta_fields`` cannot contain :class:`jax.Array` or
-      :class:`numpy.ndarray` objects.
-    data_fields: data field names. These fields *must* be JAX-compatible objects
-      such as arrays (:class:`jax.Array` or :class:`numpy.ndarray`), scalars, or
-      pytrees whose leaves are arrays or scalars. Note that ``data_fields`` may be
-      ``None``, as this is recognized by JAX as an empty pytree.
+    meta_fields: metadata field names: these are attributes which will be treated as
+      {term}`static` when this pytree is passed to :func:`jax.jit`. ``meta_fields`` is
+      optional only if ``nodetype`` is a dataclass, in which case individual fields can
+      be marked static via :func:`dataclasses.field` (see examples below).
+      Metadata fields *must* be static, hashable, immutable objects, as these objects
+      are used to generate JIT cache keys. In particular, metadata fields cannot contain
+      :class:`jax.Array` or :class:`numpy.ndarray` objects.
+    data_fields: data field names: these are attributes which will be treated as non-static
+      when this pytree is passed to :func:`jax.jit`. ``data_fields`` is optional only if
+      ``nodetype`` is a dataclass, in which case fields are assumed data fields unless
+      marked via :func:`dataclasses.field` (see examples below).
+      Data fields *must* be JAX-compatible objects such as arrays (:class:`jax.Array`
+      or :class:`numpy.ndarray`), scalars, or pytrees whose leaves are arrays or scalars.
+      Note that ``None`` is a valid data field, as JAX recognizes this as an empty pytree.
 
   Returns:
     The input class ``nodetype`` is returned unchanged after being added to JAX's
-    pytree registry. This return value allows ``register_dataclass`` to be partially
-    evaluated and used as a decorator as in the example below.
+    pytree registry, so that :func:`register_dataclass` can be used as a decorator.
 
   Examples:
+    In JAX v0.4.35 or older, you must specify ``data_fields`` and ``meta_fields``
+    in order to use this decorator:
+
+    >>> import jax
     >>> from dataclasses import dataclass
     >>> from functools import partial
-    >>>
+    ...
     >>> @partial(jax.tree_util.register_dataclass,
     ...          data_fields=['x', 'y'],
     ...          meta_fields=['op'])
@@ -976,7 +985,26 @@ def register_dataclass(
     >>> m
     MyStruct(x=Array([1., 1., 1.], dtype=float32), y=Array([0, 1, 2], dtype=int32), op='add')
 
-    Now that this class is registered, it can be used with functions in :mod:`jax.tree_util`:
+    Starting in JAX v0.4.36, the ``data_fields`` and ``meta_fields`` arguments are optional
+    for :func:`~dataclasses.dataclass` inputs, with fields defaulting to ``data_fields``
+    unless marked as static using `static` metadata in :func:`dataclasses.field`.
+
+    >>> import jax
+    >>> from dataclasses import dataclass, field
+    ...
+    >>> @jax.tree_util.register_dataclass
+    ... @dataclass
+    ... class MyStruct:
+    ...   x: jax.Array  # defaults to non-static data field
+    ...   y: jax.Array  # defaults to non-static data field
+    ...   op: str = field(metadata=dict(static=True))  # marked as static meta field.
+    ...
+    >>> m = MyStruct(x=jnp.ones(3), y=jnp.arange(3), op='add')
+    >>> m
+    MyStruct(x=Array([1., 1., 1.], dtype=float32), y=Array([0, 1, 2], dtype=int32), op='add')
+
+    Once this class is registered, it can be used with functions in :mod:`jax.tree` and
+    :mod:`jax.tree_util`:
 
     >>> leaves, treedef = jax.tree.flatten(m)
     >>> leaves
@@ -987,7 +1015,8 @@ def register_dataclass(
     MyStruct(x=Array([1., 1., 1.], dtype=float32), y=Array([0, 1, 2], dtype=int32), op='add')
 
     In particular, this registration allows ``m`` to be passed seamlessly through code
-    wrapped in :func:`jax.jit` and other JAX transformations:
+    wrapped in :func:`jax.jit` and other JAX transformations, with ``data_fields`` being
+    treated as dynamic arguments, and ``meta_fields`` being treated as static arguments:
 
     >>> @jax.jit
     ... def compiled_func(m):
@@ -999,6 +1028,21 @@ def register_dataclass(
     >>> compiled_func(m)
     Array([1., 2., 3.], dtype=float32)
   """
+  if data_fields is None or meta_fields is None:
+    if data_fields is not None or meta_fields is not None:
+      raise TypeError("register_dataclass: data_fields and meta_fields must both be specified"
+                      f" when either is specified. Got {data_fields=} {meta_fields=}.")
+    if not dataclasses.is_dataclass(nodetype):
+      raise TypeError("register_dataclass: data_fields and meta_fields are required when"
+                      f" nodetype is not a dataclass. Got {nodetype=}.")
+    data_fields = [f.name for f in dataclasses.fields(nodetype)
+                   if not f.metadata.get('static', False)]
+    meta_fields = [f.name for f in dataclasses.fields(nodetype)
+                   if f.metadata.get('static', False)]
+
+  assert meta_fields is not None
+  assert data_fields is not None
+
   # Store inputs as immutable tuples in this scope, because we close over them
   # for later evaluation. This prevents potentially confusing behavior if the
   # caller were to pass in lists that are later mutated.


Test patch:
diff --git a/tests/tree_util_test.py b/tests/tree_util_test.py
index 378e3803b..a3a8bc96e 100644
--- a/tests/tree_util_test.py
+++ b/tests/tree_util_test.py
@@ -17,7 +17,6 @@ import dataclasses
 import functools
 import pickle
 import re
-from typing import TypeVar
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -142,27 +141,6 @@ class FlatCache:
       data, meta = tree_util.tree_flatten(tree_util.tree_unflatten(meta, data))
     return FlatCache(None, leaves=data, treedef=meta)
 
-_T = TypeVar("_T")
-
-
-# Inspired by Flax.
-def pytree_node_dataclass(clz: _T, **kwargs) -> _T:
-  data_clz = dataclasses.dataclass(**kwargs)(clz)  # type: ignore
-  meta_fields = []
-  data_fields = []
-  for field_info in dataclasses.fields(data_clz):
-    is_pytree_node = field_info.metadata.get("pytree_node", True)
-    if is_pytree_node:
-      data_fields.append(field_info.name)
-    else:
-      meta_fields.append(field_info.name)
-
-  jax.tree_util.register_dataclass(
-      data_clz, data_fields, meta_fields
-  )
-
-  return data_clz
-
 
 @tree_util.register_static
 class StaticInt(int):
@@ -231,16 +209,18 @@ TREE_STRINGS = (
     "PyTreeDef(CustomNode(StaticDict[{'foo': 4, 'bar': 5}], []))",
 )
 
-@pytree_node_dataclass
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass
 class ADataclass:
   x: tuple[int, int]
   y: int
 
-@pytree_node_dataclass
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass
 class ADataclassWithMeta:
   x: tuple[int, int]
   y: int
-  z: int = dataclasses.field(metadata={"pytree_node": False})
+  z: int = dataclasses.field(metadata={"static": True})
 
 TREES += (
     (ADataclass(x=(1, 2), y=3),),
@@ -1294,6 +1274,36 @@ class TreeAliasTest(jtu.JaxTestCase):
 
 class RegistrationTest(jtu.JaxTestCase):
 
+  def test_register_dataclass_with_field_specifier(self):
+    @tree_util.register_dataclass
+    @dataclasses.dataclass
+    class Foo:
+      x: int
+      y: int = dataclasses.field(metadata=dict(static=True))
+
+    f = Foo(2, 3)
+    self.assertLen(jax.tree.leaves(f), 1)
+
+  def test_register_dataclass_field_errors(self):
+    class Foo:  # not a dataclass
+      x: int
+      y: int
+
+    msg = ("register_dataclass: data_fields and meta_fields are required"
+           " when nodetype is not a dataclass. Got nodetype=<class '.*Foo'>")
+    with self.assertRaisesRegex(TypeError, msg):
+      tree_util.register_dataclass(Foo)
+
+    msg = ("register_dataclass: data_fields and meta_fields must both be specified"\
+           r" when either is specified. Got data_fields=\['x'\] meta_fields=None.")
+    with self.assertRaisesRegex(TypeError, msg):
+      tree_util.register_dataclass(Foo, data_fields=['x'])
+
+    msg = ("register_dataclass: data_fields and meta_fields must both be specified"\
+           r" when either is specified. Got data_fields=None meta_fields=\['y'\].")
+    with self.assertRaisesRegex(TypeError, msg):
+      tree_util.register_dataclass(Foo, meta_fields=['y'])
+
   def test_register_dataclass_missing_fields(self):
     @dataclasses.dataclass
     class Foo: