Merge pull request #4056 from cool-RR:patch-2

PiperOrigin-RevId: 659565993
google · Aug 5, 2024 · 5694517 · 5694517
2 parents d20f594 + f2f76d9
commit 5694517
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/flax/struct.py b/flax/struct.py
@@ -27,8 +27,9 @@
 _T = TypeVar('_T')
 
 
-def field(pytree_node=True, **kwargs):
-  return dataclasses.field(metadata={'pytree_node': pytree_node}, **kwargs)
+def field(pytree_node=True, *, metadata=None, **kwargs):
+  return dataclasses.field(metadata=(metadata or {}) | {'pytree_node': pytree_node},
+                           **kwargs)
 
 
 @dataclass_transform(field_specifiers=(field,))  # type: ignore[literal-required]

diff --git a/tests/struct_test.py b/tests/struct_test.py
@@ -123,6 +123,12 @@ class B(A):
         class B(A, struct.PyTreeNode):
           b: int
 
+  def test_metadata_pass_through(self):
+    @struct.dataclass
+    class A:
+      foo: int = struct.field(default=9, metadata={'baz': 9})
+    assert A.__dataclass_fields__['foo'].metadata == {'baz': 9, 'pytree_node': True}
+
   @parameterized.parameters(
       {'mode': 'dataclass'},
       {'mode': 'pytreenode'},