Add pytoch 2.0 compatibility

PiperOrigin-RevId: 517483590

Author: Conchylicultor

CHANGELOG.md

@@ -23,6 +23,8 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
+*   Add `torch==2.0.0` support
+
 ## [1.4.0] - 2023-03-13
 
 *   **Add `torch` support!**

dataclass_array/array_dataclass.py

@@ -515,10 +515,18 @@ def as_xnp(self: _DcT, xnp: enp.NpModule) -> _DcT:
       return self
     # Direct `torch` <> `tf`/`jax` conversion not supported, so convert to
     # `numpy`
-    if enp.lazy.has_torch and (
-        xnp is enp.lazy.torch or self.xnp is enp.lazy.torch
-    ):
-      array_fn = lambda f: xnp.asarray(np.asarray(f.value))
+    if enp.lazy.is_torch_xnp(xnp) or enp.lazy.is_torch_xnp(self.xnp):
+
+      def _as_torch(f):
+        arr = np.asarray(f.value)
+        # Torch fail for scalar arrays:
+        # https://github.com/pytorch/pytorch/issues/97021
+        if enp.lazy.is_torch_xnp(xnp) and not arr.shape:  # Destination is torch
+          return xnp.asarray(arr.item(), dtype=lazy.as_torch_dtype(arr.dtype))
+
+        return xnp.asarray(arr)
+
+      array_fn = _as_torch
     else:
       array_fn = lambda f: xnp.asarray(f.value)
 

dataclass_array/array_dataclass_test.py

@@ -609,11 +609,9 @@ def test_torch_device():
     ]
 )
 def test_vmap(xnp: enp.NpModule):
-  import functorch
-
   vmap_fn = {
       enp.lazy.jnp: enp.lazy.jax.vmap,
-      enp.lazy.torch: functorch.vmap,
+      enp.lazy.torch: enp.lazy.torch.func.vmap,
   }[xnp]
 
   batch_shape = 3

dataclass_array/testing.py

@@ -21,6 +21,7 @@
 
 from dataclass_array import array_dataclass
 from dataclass_array.typing import FloatArray  # pylint: disable=g-multiple-import
+from etils import enp
 from etils.etree import jax as etree
 from etils.etree import Tree
 import numpy as np
@@ -61,6 +62,9 @@ def assert_allclose(
 def assert_array_equal(
     x,
     y,
+    *,
+    atol: Optional[float] = None,
+    rtol: Optional[float] = None,
 ) -> None:
   """Assert the 2 objects are equals.
 
@@ -71,9 +75,26 @@ def assert_array_equal(
   Args:
     x: First element to compare
     y: Second element to compare
+    atol: Absolute tolerance
+    rtol: Relative tolerance
   """
   assert type(x) == type(y)  # pylint: disable=unidiomatic-typecheck
   assert x.shape == y.shape
-  assert_allclose(x, y)
+  assert_allclose(x, y, atol=atol, rtol=rtol)
   if isinstance(x, array_dataclass.DataclassArray):
     assert x.xnp is y.xnp
+
+
+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 skip_torch:
+      skip = True
+  if skip:
+    import pytest  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
+
+    pytest.skip('Vectorization not supported yet with TF / Torch')

dataclass_array/vectorization.py

@@ -377,12 +377,18 @@ def _jax_vmap_cached(fn: _FnT, *, in_axes) -> _FnT:
 @functools.lru_cache(maxsize=None)
 def _torch_vmap_cached(fn: _FnT, *, in_axes) -> _FnT:
   """Like `jax.vmap` but cache the function."""
-  try:
-    import functorch  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
-  except ImportError as e:
-    epy.reraise(e, suffix='. vectorization with `pytorch` require functorch')
+  if hasattr(enp.lazy.torch, 'func'):  # torch 2.0
+    vmap = enp.lazy.torch.func.vmap
+  else:
+    try:
+      import functorch  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
+    except ImportError as e:
+      epy.reraise(
+          e, suffix='. vectorization with `pytorch<2` require functorch'
+      )
+    vmap = functorch.vmap
 
-  return functorch.vmap(
+  return vmap(
       fn,
       in_dims=in_axes,
   )

dataclass_array/vectorization_test.py

@@ -26,7 +26,6 @@
 from etils import enp
 import jax
 import pytest
-import tensorflow.experimental.numpy as tnp
 
 H = 2
 W = 3
@@ -177,16 +176,15 @@ def fn(self):
   assert a.shape == (3,)
   assert a.x == 5
 
-  # Vectorization supported
-  if xnp not in [
-      tnp,
-  ]:
-    a = a.fn()
-  assert a.xnp is xnp
+  # Tree-map supported
+  a = jax.tree_util.tree_map(lambda x: x, a)
   assert a.shape == (3,)
   assert a.x == 5
 
-  # Tree-map supported
-  a = jax.tree_util.tree_map(lambda x: x, a)
+  # Vectorization supported
+  dca.testing.skip_vmap_unavailable(xnp)
+
+  a = a.fn()
+  assert a.xnp is xnp
   assert a.shape == (3,)
   assert a.x == 5