Repository: google/flax

Issue:
[struct.dataclass] Consider adding optional `kw_only` arguments
I often run into the following issue:

```python
from flax import struct

class Foo(struct.PyTreeNode):
  bar: int = struct.field(pytree_node=False, default=1)

class Baz(Foo):
  qux: str
```

Since `qux` does not have a default value, I get:
```
Fields without default values cannot appear after fields with default values
```

Can we consider adding a simple wrapper to `dataclasses.dataclass(kw_only=True)`?

It should be easy for the `struct.dataclass`, we can maybe have another object for inheritance, like `PyTreeNodeKwOnly`?


Discussion:
#3645 allows you to pass in general kwargs to `struct.dataclass`, so you should be able to use `kw_only=True` now.

#3651 adds a test to test this `kw_only` functionality. Please note that `kw_only` is only available for python 3.10 or higher. Currently Flax supports, python 3.9 or higher, and so this test is commented out until Flax upgrades its minimum version requirement.
Thanks @chiamp, would #3645 pass the test case in the issue above?
Yes, you could do something like this:
```
from flax import struct
import functools

@functools.partial(struct.dataclass, kw_only=True)
class Foo():
  bar: int = struct.field(pytree_node=False, default=1)

@struct.dataclass
class Baz(Foo):
  qux: str

b = Baz('qux_value', bar=123)
print(b)
# Baz(bar=123, qux='qux_value')
```
Thanks, @chiamp, that is a different test case, though.
The original one still fails.

In the second test case, Pylance also does not recognise that `Baz` has a `bar` attribute: `No parameter named "bar"`.

Is there any plan to fully support type checking with this and propagate the change to the `PyTreeNode` class?


Resolving patch:
diff --git a/flax/struct.py b/flax/struct.py
--- a/flax/struct.py
+++ b/flax/struct.py
@@ -227,8 +227,8 @@ class PyTreeNode:
     >>> model_grad = jax.grad(loss_fn)(model)
   """
 
-  def __init_subclass__(cls):
-    dataclass(cls)  # pytype: disable=wrong-arg-types
+  def __init_subclass__(cls, **kwargs):
+    dataclass(cls, **kwargs)  # pytype: disable=wrong-arg-types
 
   def __init__(self, *args, **kwargs):
     # stub for pytype


Test patch:
diff --git a/tests/struct_test.py b/tests/struct_test.py
--- a/tests/struct_test.py
+++ b/tests/struct_test.py
@@ -18,7 +18,7 @@
 from typing import Any
 
 import jax
-from absl.testing import absltest
+from absl.testing import absltest, parameterized
 from jax._src.tree_util import prefix_errors
 
 from flax import struct
@@ -34,7 +34,7 @@ class Point:
   meta: Any = struct.field(pytree_node=False)
 
 
-class StructTest(absltest.TestCase):
+class StructTest(parameterized.TestCase):
   def test_no_extra_fields(self):
     p = Point(x=1, y=2, meta={})
     with self.assertRaises(dataclasses.FrozenInstanceError):
@@ -93,24 +93,68 @@ class A(struct.PyTreeNode):
       a: int
 
   # TODO(marcuschiam): Uncomment when Flax upgrades to Python 3.10.
-  # def test_kw_only(self):
-  #   @struct.dataclass
-  #   class A:
-  #     a: int = 1
-
-  #   with self.assertRaisesRegex(TypeError, "non-default argument 'b' follows default argument"):
+  # @parameterized.parameters(
+  #     {'mode': 'dataclass'},
+  #     {'mode': 'pytreenode'},
+  # )
+  # def test_kw_only(self, mode):
+  #   if mode == 'dataclass':
   #     @struct.dataclass
+  #     class A:
+  #       a: int = 1
+
+  #     @functools.partial(struct.dataclass, kw_only=True)
   #     class B(A):
   #       b: int
+  #   elif mode == 'pytreenode':
+  #     class A(struct.PyTreeNode):
+  #       a: int = 1
 
-  #   @functools.partial(struct.dataclass, kw_only=True)
-  #   class B(A):
-  #     b: int
+  #     class B(A, struct.PyTreeNode, kw_only=True):
+  #       b: int
 
   #   obj = B(b=2)
   #   self.assertEqual(obj.a, 1)
   #   self.assertEqual(obj.b, 2)
 
+  #   with self.assertRaisesRegex(TypeError, "non-default argument 'b' follows default argument"):
+  #     if mode == 'dataclass':
+  #       @struct.dataclass
+  #       class B(A):
+  #         b: int
+  #     elif mode == 'pytreenode':
+  #       class B(A, struct.PyTreeNode):
+  #         b: int
+
+  # TODO(marcuschiam): Uncomment when Flax upgrades to Python 3.10.
+  # @parameterized.parameters(
+  #     {'mode': 'dataclass'},
+  #     {'mode': 'pytreenode'},
+  # )
+  # def test_mutable(self, mode):
+  #   if mode == 'dataclass':
+  #     @struct.dataclass
+  #     class A:
+  #       a: int = 1
+
+  #     @functools.partial(struct.dataclass, frozen=False)
+  #     class B:
+  #       b: int = 1
+  #   elif mode == 'pytreenode':
+  #     class A(struct.PyTreeNode):
+  #       a: int = 1
+
+  #     class B(struct.PyTreeNode, frozen=False):
+  #       b: int = 1
+
+  #   obj = A()
+  #   with self.assertRaisesRegex(dataclasses.FrozenInstanceError, "cannot assign to field 'a'"):
+  #     obj.a = 2
+
+  #   obj = B()
+  #   obj.b = 2
+  #   self.assertEqual(obj.b, 2)
+
 
 if __name__ == '__main__':
   absltest.main()