Add pytorch support to dataclass_array

PiperOrigin-RevId: 509333699

Author: Conchylicultor

CHANGELOG.md

@@ -23,6 +23,9 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
+*   Add `torch` support (experimental). Require to call
+    `dca.activate_torch_support()`
+
 ## [1.3.0] - 2023-01-16
 
 *   Added: Support for static `dca.DataclassArray` (dataclasses with only

dataclass_array/__init__.py

@@ -37,6 +37,9 @@
 from dataclass_array.ops import stack
 from dataclass_array.vectorization import vectorize_method
 
+# TODO(epot): Remove once Torch has better numpy API
+from etils.enp import activate_torch_support
+
 # `dca.testing` do not depend on pytest or other heavy deps, so is safe to
 # import
 from dataclass_array import testing

dataclass_array/array_dataclass.py

@@ -195,8 +195,9 @@ def __init_subclass__(cls, **kwargs):
     super().__init_subclass__(**kwargs)
     # TODO(epot): Could have smart __repr__ which display types if array have
     # too many values (maybe directly in `edc.field(repr=...)`).
-    edc.dataclass(kw_only=True, repr=True)(cls)
-    cls._dca_tree_map_registered = False
+    edc.dataclass(kw_only=True, repr=True, auto_cast=False)(cls)
+    cls._dca_jax_tree_registered = False
+    cls._dca_torch_tree_registered = False
     # Typing annotations have to be lazily evaluated (to support
     # `from __future__ import annotations` and forward reference)
     # To avoid costly `typing.get_type_hints` which perform `eval` and `str`
@@ -217,10 +218,20 @@ def __post_init__(self) -> None:
       _init_cls(self)
 
     # Register the tree_map here instead of `__init_subclass__` as `jax` may
-    # not have been registered yet during import
-    if enp.lazy.has_jax and not cls._dca_tree_map_registered:  # pylint: disable=protected-access
+    # not have been imported yet during import.
+    if enp.lazy.has_jax and not cls._dca_jax_tree_registered:  # pylint: disable=protected-access
       enp.lazy.jax.tree_util.register_pytree_node_class(cls)
-      cls._dca_tree_map_registered = True  # pylint: disable=protected-access
+      cls._dca_jax_tree_registered = True  # pylint: disable=protected-access
+
+    if enp.lazy.has_torch and not cls._dca_torch_tree_registered:  # pylint: disable=protected-access
+      # Note: Torch is updating it's tree API to make it public and use `optree`
+      # as backend: https://github.com/pytorch/pytorch/issues/65761
+      enp.lazy.torch.utils._pytree._register_pytree_node(  # pylint: disable=protected-access
+          cls,
+          flatten_fn=lambda a: a.tree_flatten(),
+          unflatten_fn=lambda vals, ctx: cls.tree_unflatten(ctx, vals),
+      )
+      cls._dca_torch_tree_registered = True  # pylint: disable=protected-access
 
     # Validate and normalize array fields
     # * Maybe cast (list, np) -> xnp
@@ -442,14 +453,28 @@ def as_tf(self: _DcT) -> _DcT:
     """Returns the instance as containing `tf.Tensor`."""
     return self.as_xnp(enp.lazy.tnp)
 
+  def as_torch(self: _DcT) -> _DcT:
+    """Returns the instance as containing `torch.Tensor`."""
+    return self.as_xnp(enp.lazy.torch)
+
   def as_xnp(self: _DcT, xnp: enp.NpModule) -> _DcT:
     """Returns the instance as containing `xnp.ndarray`."""
     if xnp is self.xnp:  # No-op
       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))
+    else:
+      array_fn = lambda f: xnp.asarray(f.value)
 
     # Update all childs
     new_self = self._map_field(
-        array_fn=lambda f: xnp.asarray(f.value),
+        array_fn=array_fn,
         dc_fn=lambda f: f.value.as_xnp(xnp),
     )
     return new_self
@@ -518,6 +543,17 @@ def _get_xnp(f: _ArrayField) -> enp.NpModule:
       return None
     xnp = _infer_xnp(xnps)
 
+    if (
+        enp.lazy.has_torch
+        and xnp is enp.lazy.torch
+        and not hasattr(enp.lazy.torch, '__etils_np_mode__')
+    ):
+      raise ValueError(
+          'torch support currently require to call:\n'
+          'import dataclass_array as dca\n'
+          'dca.activate_torch_support()'
+      )
+
     def _cast_field(f: _ArrayField) -> None:
       try:
         new_value = np_utils.asarray(

dataclass_array/array_dataclass_test.py

@@ -28,7 +28,7 @@
 import tensorflow as tf
 
 # Activate the fixture
-set_tnp = enp.testing.set_tnp
+enable_torch_tf_np_mode = enp.testing.enable_torch_tf_np_mode
 
 # TODO(epot): Test dtype `complex`, `str`
 
@@ -70,8 +70,8 @@ def assert_val(p: Point, shape: Shape, xnp: enp.NpModule = None):
     _assert_common(p, shape=shape, xnp=xnp)
     assert p.x.shape == shape
     assert p.y.shape == shape
-    assert p.x.dtype == np.float32
-    assert p.y.dtype == np.float32
+    assert enp.lazy.as_dtype(p.x.dtype) == np.float32
+    assert enp.lazy.as_dtype(p.y.dtype) == np.float32
     assert isinstance(p.x, xnp.ndarray)
     assert isinstance(p.y, xnp.ndarray)
 
@@ -98,8 +98,8 @@ def assert_val(p: Isometrie, shape: Shape, xnp: enp.NpModule = None):
     _assert_common(p, shape=shape, xnp=xnp)
     assert p.r.shape == shape + (3, 3)
     assert p.t.shape == shape + (2,)
-    assert p.r.dtype == np.float32
-    assert p.t.dtype == np.int32
+    assert enp.lazy.as_dtype(p.r.dtype) == np.float32
+    assert enp.lazy.as_dtype(p.t.dtype) == np.int32
     assert isinstance(p.r, xnp.ndarray)
     assert isinstance(p.t, xnp.ndarray)
 
@@ -226,8 +226,8 @@ def assert_val(p: WithStatic, shape: Shape, xnp: enp.NpModule = None):
     NestedOnlyStatic.assert_val(p.nested_static, shape, xnp=xnp)
     assert p.x.shape == shape + (3,)
     assert p.y.shape == shape + (2, 2)
-    assert p.x.dtype == np.float32
-    assert p.y.dtype == np.float32
+    assert enp.lazy.as_dtype(p.x.dtype) == np.float32
+    assert enp.lazy.as_dtype(p.y.dtype) == np.float32
     assert isinstance(p.x, xnp.ndarray)
     assert isinstance(p.y, xnp.ndarray)
     # Static field is correctly forwarded
@@ -546,12 +546,15 @@ def test_convert(
 ):
   p = dca_cls.make(xnp=xnp, shape=(2,))
   assert p.xnp is xnp
+
   assert p.as_np().xnp is enp.lazy.np
   assert p.as_jax().xnp is enp.lazy.jnp
   assert p.as_tf().xnp is enp.lazy.tnp
+  assert p.as_torch().xnp is enp.lazy.torch
   assert p.as_xnp(np).xnp is enp.lazy.np
   assert p.as_xnp(enp.lazy.jnp).xnp is enp.lazy.jnp
   assert p.as_xnp(enp.lazy.tnp).xnp is enp.lazy.tnp
+  assert p.as_xnp(enp.lazy.torch).xnp is enp.lazy.torch
   # Make sure the nested class are also updated
   dca_cls.assert_val(p.as_jax(), (2,), xnp=enp.lazy.jnp)
 
@@ -587,24 +590,44 @@ def test_infer_np(xnp: enp.NpModule):
 
 
 @parametrize_dataclass_arrays
-def test_jax_tree_map(dca_cls: DcaTest):
+@pytest.mark.parametrize(
+    'tree_map',
+    [
+        enp.lazy.jax.tree_map,
+        enp.lazy.torch.utils._pytree.tree_map,
+    ],
+)
+def test_torch_tree_map(tree_map, dca_cls: DcaTest):
   p = dca_cls.make(shape=(3,), xnp=np)
-  p = enp.lazy.jax.tree_map(lambda x: x[None, ...], p)
+  p = tree_map(lambda x: x[None, ...], p)
   dca_cls.assert_val(p, (1, 3), xnp=np)
 
 
-def test_jax_vmap():
+@enp.testing.parametrize_xnp(
+    restrict=[
+        'jnp',
+        'torch',
+    ]
+)
+def test_vmap(xnp: enp.NpModule):
+  import functorch
+
+  vmap_fn = {
+      enp.lazy.jnp: enp.lazy.jax.vmap,
+      enp.lazy.torch: functorch.vmap,
+  }[xnp]
+
   batch_shape = 3
 
-  @enp.lazy.jax.vmap
+  @vmap_fn
   def fn(p: WithStatic) -> WithStatic:
     assert isinstance(p, WithStatic)
     assert p.shape == ()  # pylint:disable=g-explicit-bool-comparison
     return p.replace(x=p.x + 1)
 
-  x = WithStatic.make((batch_shape,), xnp=enp.lazy.jnp)
+  x = WithStatic.make((batch_shape,), xnp=xnp)
   y = fn(x)
-  WithStatic.assert_val(y, (batch_shape,), xnp=enp.lazy.jnp)
+  WithStatic.assert_val(y, (batch_shape,), xnp=xnp)
   # pos was updated
   np.testing.assert_allclose(y.x, np.ones((batch_shape, 3)))
   np.testing.assert_allclose(y.y, np.zeros((batch_shape, 2, 2)))
@@ -628,8 +651,8 @@ class PointNoCast(dca.DataclassArray):
       y=xnp.array([1, 2, 3], dtype=np.uint8),
   )
   assert p.shape == (3,)
-  assert p.x.dtype == np.float16
-  assert p.y.dtype == np.uint8
+  assert enp.lazy.as_dtype(p.x.dtype) == np.float16
+  assert enp.lazy.as_dtype(p.y.dtype) == np.uint8
 
 
 @enp.testing.parametrize_xnp()
@@ -689,7 +712,13 @@ class PointDynamicShape(dca.DataclassArray):
   assert dca.stack([p, p]).shape == (2,) + batch_shape
 
   # Incompatible shape will raise an error
-  with pytest.raises((ValueError, tf.errors.InvalidArgumentError)):
+  expected_exception_cls = {
+      enp.lazy.np: ValueError,
+      enp.lazy.jnp: ValueError,
+      enp.lazy.tnp: tf.errors.InvalidArgumentError,
+      enp.lazy.torch: RuntimeError,
+  }
+  with pytest.raises(expected_exception_cls[xnp]):
     dca.stack([p, p2])
 
   if batch_shape:

dataclass_array/import_test.py

@@ -20,12 +20,20 @@
 
 from __future__ import annotations
 
+import dataclasses
 import sys
 
 import dataclass_array as dca
+from etils import enp
+import pytest
 
-del dca
+
+@dataclasses.dataclass(frozen=True)
+class A(dca.DataclassArray):
+  x: dca.typing.f32['*s']
 
 
 def test_lazy():
-  pass
+
+  x = A(x=[1.0, 2.0])
+  assert x.xnp is enp.lazy.np

dataclass_array/utils/np_utils_test.py

@@ -22,6 +22,9 @@
 import numpy as np
 import pytest
 
+# Activate the fixture
+enable_torch_tf_np_mode = enp.testing.enable_torch_tf_np_mode
+
 
 @enp.testing.parametrize_xnp()
 def test_get_xnp(xnp: enp.NpModule):

dataclass_array/vectorization.py

@@ -311,9 +311,19 @@ def _vmap_method(
   if xnp is enp.lazy.np:
     return _vmap_method_np(args, map_non_static=map_non_static)
   elif xnp is enp.lazy.jnp:
-    return _vmap_method_jnp(args, map_non_static=map_non_static)
+    return _vmap_method_jax_torch(
+        args,
+        map_non_static=map_non_static,
+        make_vmap_fn=_jax_vmap_cached,
+    )
   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,
+    )
   raise TypeError(f'Invalid numpy module: {xnp}')
 
 
@@ -334,10 +344,11 @@ def _vmap_method_np(
   return tree_utils.tree_map(_stack, *outs)
 
 
-def _vmap_method_jnp(
+def _vmap_method_jax_torch(
     args: inspect_utils.BoundArgs[Any, _OutT],
     *,
     map_non_static: _MapNonStatic,
+    make_vmap_fn: Any,
 ) -> _OutT:
   """vectorization using `jax` backend."""
 
@@ -349,21 +360,35 @@ def _vmap_method_jnp(
   in_axes = tuple(arg.value for arg in in_axes_args)
 
   # Vectorize self and args
-  vfn = _vmap_cached(args.fn, in_axes=in_axes)
+  vfn = make_vmap_fn(args.fn, in_axes=in_axes)
 
   # Call `vfn(self, *args, **kwargs)`
   return args.call(vfn)
 
 
 @functools.lru_cache(maxsize=None)
-def _vmap_cached(fn: _FnT, *, in_axes) -> _FnT:
+def _jax_vmap_cached(fn: _FnT, *, in_axes) -> _FnT:
   """Like `jax.vmap` but cache the function."""
   return enp.lazy.jax.vmap(
       fn,
       in_axes=in_axes,
   )
 
 
+@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')
+
+  return functorch.vmap(
+      fn,
+      in_dims=in_axes,
+  )
+
+
 def _vmap_method_tf(
     args: inspect_utils.BoundArgs[Any, _OutT],
     *,

dataclass_array/vectorization_test.py

@@ -34,7 +34,7 @@
 X1 = 5
 
 # Activate the fixture
-set_tnp = enp.testing.set_tnp
+enable_torch_tf_np_mode = enp.testing.enable_torch_tf_np_mode
 
 
 @pytest.mark.parametrize(
@@ -179,7 +179,9 @@ def fn(self):
   assert a.x == 5
 
   # Vectorization supported
-  if xnp != tnp:
+  if xnp not in [
+      tnp,
+  ]:
     a = a.fn()
   assert a.xnp is xnp
   assert a.shape == (3,)

pyproject.toml

@@ -43,6 +43,7 @@ dev = [
     "chex",
     "jax[cpu]",
     "tf-nightly",
+    "torch",
 ]
 
 [tool.pyink]