Activate test for TF and dca.vectorize_method support

PiperOrigin-RevId: 561646634

Author: Etienne Pot

CHANGELOG.md

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

dataclass_array/testing.py

@@ -88,10 +88,7 @@ def assert_array_equal(
 def skip_vmap_unavailable(xnp: enp.NpModule, *, skip_torch: str = '') -> None:
   """Skip the test when vmap not available."""
   skip = False
-  if enp.lazy.is_tf_xnp(xnp):
-    # TODO(b/152678472): TF do not support vmap & tf.nest
-    skip = True
-  elif enp.lazy.is_torch_xnp(xnp):
+  if enp.lazy.is_torch_xnp(xnp):
     if skip_torch:
       skip = True
   if skip:

dataclass_array/vectorization.py

@@ -325,13 +325,8 @@ def _vmap_method(
         make_vmap_fn=make_vmap_fn,
     )
   elif enp.lazy.is_tf_xnp(xnp):
-    # return _vmap_method_tf(args, map_non_static=map_non_static)
+    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}')