support passing arguments directly to the struct.dataclass decorator

PiperOrigin-RevId: 694517927

Author: Qwlouse

flax/struct.py

@@ -14,8 +14,10 @@
 
 """Utilities for defining custom classes that can be used with jax transformations."""
 
+from collections.abc import Callable
 import dataclasses
-from typing import TypeVar
+import functools
+from typing import TypeVar, overload
 
 import jax
 from typing_extensions import (
@@ -33,7 +35,22 @@ def field(pytree_node=True, *, metadata=None, **kwargs):
 
 
 @dataclass_transform(field_specifiers=(field,))  # type: ignore[literal-required]
+@overload
 def dataclass(clz: _T, **kwargs) -> _T:
+  ...
+
+
+@dataclass_transform(field_specifiers=(field,))  # type: ignore[literal-required]
+@overload
+def dataclass(**kwargs) -> Callable[[_T], _T]:
+  ...
+
+
+@dataclass_transform(field_specifiers=(field,))  # type: ignore[literal-required]
+def dataclass(
+    clz: _T | None = None,
+    **kwargs,
+) -> _T | Callable[[_T], _T]:
   """Create a class which can be passed to functional transformations.
 
   .. note::
@@ -99,9 +116,15 @@ class method that provides the smart constructor.
 
   Args:
     clz: the class that will be transformed by the decorator.
+    **kwargs: arguments to pass to the dataclass constructor.
+
   Returns:
     The new class.
   """
+  # Support passing arguments to the decorator (e.g. @dataclass(kw_only=True))
+  if clz is None:
+    return functools.partial(dataclass, **kwargs)
+
   # check if already a flax dataclass
   if '_flax_dataclass' in clz.__dict__:
     return clz
@@ -119,7 +142,7 @@ class method that provides the smart constructor.
       meta_fields.append(field_info.name)
 
   def replace(self, **updates):
-    """ "Returns a new object replacing the specified fields with new values."""
+    """Returns a new object replacing the specified fields with new values."""
     return dataclasses.replace(self, **updates)
 
   data_clz.replace = replace

tests/struct_test.py

@@ -15,7 +15,6 @@
 """Tests for flax.struct."""
 
 import dataclasses
-import functools
 from typing import Any
 
 import jax
@@ -49,8 +48,8 @@ def test_mutation(self):
       p.y = 3
 
   def test_slots(self):
-    slots_dataclass = functools.partial(struct.dataclass, frozen=False, slots=True)
-    @slots_dataclass
+
+    @struct.dataclass(frozen=False, slots=True)
     class SlotsPoint:
       x: float
       y: float
@@ -100,7 +99,7 @@ def test_kw_only(self, mode):
       class A:
         a: int = 1
 
-      @functools.partial(struct.dataclass, kw_only=True)
+      @struct.dataclass(kw_only=True)
       class B(A):
         b: int
     elif mode == 'pytreenode':
@@ -139,7 +138,7 @@ def test_mutable(self, mode):
       class A:
         a: int = 1
 
-      @functools.partial(struct.dataclass, frozen=False)
+      @struct.dataclass(frozen=False)
       class B:
         b: int = 1
     elif mode == 'pytreenode':