Add tf.nest support to dca

PiperOrigin-RevId: 549246873

Author: Conchylicultor

CHANGELOG.md

@@ -23,6 +23,8 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
+*   Add `tf.nest` compatibility for `DataclassArray`.
+
 ## [1.4.2] - 2023-07-10
 
 *   Add `dca.concat` method in addition to `dca.stack`.

dataclass_array/array_dataclass.py

@@ -581,6 +581,24 @@ def cuda(self: _DcT, *args, **kwargs) -> _DcT:
 
   # ====== Internal ======
 
+  @epy.cached_property
+  def _all_fields_empty(self) -> bool:
+    """Returns True if the `dataclass_array` is invalid."""
+    if not self._array_fields:  # All fields are `None` / `object`
+      # No fields have been defined.
+      # This can be the case internally by jax which apply some
+      # `tree_map(lambda x: sentinel)`.
+      return True
+
+    # `tf.nest` sometimes replace values by dummy `.` inside
+    # `assert_same_structure`
+    if enp.lazy.has_tf:
+      from tensorflow.python.util import nest_util  # pytype: disable=import-error  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+
+      if any(f.value is nest_util._DOT for f in self._array_fields):  # pylint: disable=protected-access,not-an-iterable
+        return True
+    return False
+
   @epy.cached_property
   def _all_array_fields(self) -> dict[str, _ArrayField]:
     """All array fields, including `None` values."""
@@ -603,7 +621,7 @@ def _array_fields(self) -> list[_ArrayField]:
 
   def _cast_xnp_dtype_inplace(self) -> Optional[enp.NpModule]:
     """Validate `xnp` are consistent and cast `np` -> `xnp` in-place."""
-    if not self._array_fields:  # All fields are `None` / `object`
+    if self._all_fields_empty:  # pylint: disable=using-constant-test
       return None
 
     # Validate the dtype
@@ -627,12 +645,19 @@ def _get_xnp(f: _ArrayField) -> enp.NpModule:
 
     def _cast_field(f: _ArrayField) -> None:
       try:
-        new_value = np_utils.asarray(
-            f.value,
-            xnp=xnp,
-            dtype=f.dtype,
-            cast_dtype=self.__dca_params__.cast_dtype,
-        )
+        # Supports for TensorSpec (e.g. in `tf.function` signature)
+        if enp.lazy.is_tf_xnp(xnp) and isinstance(
+            f.value, enp.lazy.tf.TensorSpec
+        ):
+          # TODO(epot): Actually check the dtype
+          new_value = f.value
+        else:
+          new_value = np_utils.asarray(
+              f.value,
+              xnp=xnp,
+              dtype=f.dtype,
+              cast_dtype=self.__dca_params__.cast_dtype,
+          )
         self._setattr(f.name, new_value)
         # After the field has been set, we validate the shape
         f.assert_shape()
@@ -648,9 +673,7 @@ def _cast_field(f: _ArrayField) -> None:
 
   def _broadcast_shape_inplace(self) -> Optional[Shape]:
     """Validate the shapes are consistent and broadcast values in-place."""
-    if not self._array_fields:  # No fields have been defined.
-      # This can be the case internally by jax which apply some
-      # `tree_map(lambda x: sentinel)`.
+    if self._all_fields_empty:  # pylint: disable=using-constant-test
       return None
 
     # First collect all shapes and compute the final shape.
@@ -735,7 +758,7 @@ def _apply_field_dn(f: _ArrayField):
       else:
         return array_fn(f)
 
-    new_values = {f.name: _apply_field_dn(f) for f in self._array_fields}  # pylint: disable=not-an-iterable
+    new_values = {f.name: _apply_field_dn(f) for f in self._array_fields}  # pylint: disable=not-an-iterable,protected-access
     # For performance, do not call replace to save the constructor call
     if not _inplace:
       return self.replace(**new_values)
@@ -792,6 +815,18 @@ def tree_unflatten(
         self._setattr(k, v)  # pylint: disable=protected-access
     return self
 
+  def __tf_flatten__(self) -> tuple[_TreeMetadata, list[DcOrArray]]:
+    components, metadata = self.tree_flatten()
+    return metadata, components
+
+  @classmethod
+  def __tf_unflatten__(
+      cls: Type[_DcT],
+      metadata: _TreeMetadata,
+      components: list[DcOrArray],
+  ) -> _DcT:
+    return cls.tree_unflatten(metadata, components)
+
   def _setattr(self, name: str, value: Any) -> None:
     """Like setattr, but support `frozen` dataclasses."""
     object.__setattr__(self, name, value)

dataclass_array/array_dataclass_test.py

@@ -20,7 +20,7 @@
 
 import dataclass_array as dca
 from dataclass_array.typing import FloatArray, IntArray, f32, i32  # pylint: disable=g-multiple-import
-from dataclass_array.typing import Shape  # pylint: disable=g-multiple-import
+from dataclass_array.typing import Shape  # pylint: disable=g-multiple-import,g-importing-member
 from etils import enp
 import numpy as np
 import pytest
@@ -585,6 +585,7 @@ def test_infer_np(xnp: enp.NpModule):
 @pytest.mark.parametrize(
     'tree_map',
     [
+        enp.lazy.tf.nest.map_structure,
         enp.lazy.jax.tree_map,
         enp.lazy.torch.utils._pytree.tree_map,
     ],
@@ -595,6 +596,17 @@ 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))
+
+  for ex in ds:
+    assert isinstance(ex, Point)
+    assert ex.x.shape == ()  # pylint: disable=g-explicit-bool-comparison
+    assert ex.y.shape == ()  # pylint: disable=g-explicit-bool-comparison
+
+
 def test_torch_device():
   p = Nested.make(shape=(2,), xnp=enp.lazy.torch)
   p = p.cpu()

dataclass_array/utils/inspect_utils.py

@@ -50,7 +50,6 @@ class Signature(Generic[_OutT]):
 
   y = bound_args.call()  # Call the function
   ```
-
   """
 
   fn: _Fn[_OutT]
@@ -202,7 +201,7 @@ def map_bound_arg(
       self,
       fn: Callable[[BoundArg[_ArgT]], _NewArgT],
   ) -> BoundArgs[_NewArgT, _OutT]:
-    """Apply validation/modification to the arguments value."""
+    """Apply validation/modification to all `BoundArg`."""
 
     def _fn(arg: BoundArg[_ArgT]) -> _NewArgT:  # pytype: disable=invalid-annotation
       try:
@@ -213,9 +212,14 @@ def _fn(arg: BoundArg[_ArgT]) -> _NewArgT:  # pytype: disable=invalid-annotation
             prefix=f'For arg {arg.pos} ({arg.name!r}):\n',
         )
 
+    return self.replace_args_values({arg.name: _fn(arg) for arg in self})
+
+  def replace_args_values(
+      self, new_values: dict[str, _NewArgT]
+  ) -> BoundArgs[_NewArgT, _OutT]:
     bound_args = inspect.BoundArguments(
         signature=self.bound_args.signature,
-        arguments={arg.name: _fn(arg) for arg in self},  # pytype: disable=wrong-arg-types
+        arguments=new_values,  # pytype: disable=wrong-arg-types
     )
     return self.replace(bound_args=bound_args)  # pytype: disable=attribute-error
 

dataclass_array/vectorization.py

@@ -22,7 +22,7 @@
 
 from dataclass_array import array_dataclass
 from dataclass_array import ops
-from dataclass_array.typing import DcOrArray, Shape  # pylint: disable=g-multiple-import
+from dataclass_array.typing import DcOrArray, Shape  # pylint: disable=g-multiple-import,g-importing-member
 from dataclass_array.utils import inspect_utils
 from dataclass_array.utils import np_utils
 from dataclass_array.utils import py_utils
@@ -307,21 +307,30 @@ def _vmap_method(
     xnp: enp.NpModule,
 ) -> _Out:
   """Vectorize self using the `xnp` backend. Assume `self` was flatten."""
-  if xnp is enp.lazy.np:
+  is_jax = enp.lazy.is_jax_xnp(xnp)
+  is_torch = enp.lazy.is_torch_xnp(xnp)
+
+  if enp.lazy.is_np_xnp(xnp):
     return _vmap_method_np(args, map_non_static=map_non_static)
-  elif xnp is enp.lazy.jnp:
+  elif is_jax or is_torch:
+    if is_jax:
+      make_vmap_fn = _jax_vmap_cached
+    elif is_torch:
+      make_vmap_fn = _torch_vmap_cached
+    else:
+      raise ValueError('Unexpected')
     return _vmap_method_jax_torch(
         args,
         map_non_static=map_non_static,
-        make_vmap_fn=_jax_vmap_cached,
+        make_vmap_fn=make_vmap_fn,
     )
-  elif xnp is enp.lazy.tnp:
-    return _vmap_method_tf(args, map_non_static=map_non_static)
-  elif xnp is enp.lazy.torch:
-    return _vmap_method_jax_torch(
-        args,
-        map_non_static=map_non_static,
-        make_vmap_fn=_torch_vmap_cached,
+  elif enp.lazy.is_tf_xnp(xnp):
+    # return _vmap_method_tf(args, map_non_static=map_non_static)
+
+    # TODO(epot): Use `tf.vectorized_map()` once TF support custom nesting
+    raise NotImplementedError(
+        'vectorization not supported in TF yet due to lack of `tf.nest` '
+        'support. Please upvote or comment b/152678472.'
     )
   raise TypeError(f'Invalid numpy module: {xnp}')
 
@@ -400,13 +409,52 @@ def _vmap_method_tf(
     map_non_static: _MapNonStatic,
 ) -> _OutT:
   """vectorization using `tf` backend."""
-  # TODO(epot): Use `tf.vectorized_map()` once TF support custom nesting
-  raise NotImplementedError(
-      'vectorization not supported in TF yet due to lack of `tf.nest` '
-      'support. Please upvote or comment b/152678472.'
+
+  # Flatten args
+
+  args_info = args.map(lambda _: None)
+  # ... except the non-static ones
+  args_info = map_non_static(lambda _: 0, args_info)
+
+  # Split args in static/non-static
+  static_args = {}
+  nonstatic_args = {}
+  for a, ai in zip(args, args_info):
+    assert a.name == ai.name
+    if ai.value is None:
+      static_args[a.name] = a.value
+    else:
+      nonstatic_args[a.name] = a.value
+
+  def new_fn(non_statics, statics):
+    # Merge args and call the function
+    new_args = args.replace_args_values(dict(**non_statics, **statics))
+    return new_args.call()
+
+  # `vectorized_map(` uses autograph, which fails, so use tf.map_fn instead
+  return _better_map_fn(  #
+      functools.partial(new_fn, statics=static_args),
+      nonstatic_args,
   )
 
 
+# tf.map_fn do not support different output signature:
+def _better_map_fn(fn, elems, **kwargs):
+  """Like `tf.map_fn`."""
+  tf = enp.lazy.tf
+  if 'fn_output_signature' not in kwargs:
+    elem_spec = tf.nest.map_structure(
+        lambda t: tf.type_spec_from_value(t)._unbatch(), elems  # pylint: disable=protected-access
+    )
+    output_spec = tf.nest.map_structure(
+        tf.type_spec_from_value,
+        tf.function(fn).get_concrete_function(elem_spec).structured_outputs,
+    )
+    kwargs['fn_output_signature'] = output_spec
+
+  return tf.map_fn(fn, elems, **kwargs)
+
+
 def _stack(*vals: _OutT) -> _OutT:
   """Stack the given tree."""
   assert vals