Fix tf.data compatibility for DataclassArray

PiperOrigin-RevId: 549577729

Author: Conchylicultor

CHANGELOG.md

@@ -23,7 +23,7 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
-*   Add `tf.nest` compatibility for `DataclassArray`.
+*   Add `tf.nest`/`tf.data` compatibility for `DataclassArray`.
 
 ## [1.4.2] - 2023-07-10
 

dataclass_array/array_dataclass.py

@@ -765,10 +765,10 @@ def _apply_field_dn(f: _ArrayField):
     else:
       return self
 
-  def tree_flatten(self) -> tuple[list[DcOrArray], _TreeMetadata]:
+  def tree_flatten(self) -> tuple[tuple[DcOrArray, ...], _TreeMetadata]:
     """`jax.tree_utils` support."""
     # We flatten all values (and not just the non-None ones)
-    array_field_values = [f.value for f in self._all_array_fields.values()]
+    array_field_values = tuple(f.value for f in self._all_array_fields.values())
     metadata = _TreeMetadata(
         array_field_names=list(self._all_array_fields.keys()),
         non_array_field_kwargs={
@@ -815,7 +815,7 @@ def tree_unflatten(
         self._setattr(k, v)  # pylint: disable=protected-access
     return self
 
-  def __tf_flatten__(self) -> tuple[_TreeMetadata, list[DcOrArray]]:
+  def __tf_flatten__(self) -> tuple[_TreeMetadata, tuple[DcOrArray, ...]]:
     components, metadata = self.tree_flatten()
     return metadata, components
 

dataclass_array/array_dataclass_test.py

@@ -596,7 +596,6 @@ def test_tree_map(tree_map, dca_cls: DcaTest):
   dca_cls.assert_val(p, (1, 3), xnp=np)
 
 
-@pytest.mark.skip('tf.data fail currently')  # TODO(epot): Restore
 def test_tf_data():
   ds = tf.data.Dataset.range(3)
   ds = ds.map(lambda x: Point(x=x, y=x))