Add methods for torch device

PiperOrigin-RevId: 515367102

Author: Conchylicultor

CHANGELOG.md

@@ -24,6 +24,8 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 ## [Unreleased]
 
 *   **Add `torch` support!**
+*   Add `.cpu()`, `.cuda()`, `.to()` methods to move the dataclass from
+    devices when using torch.
 
 ## [1.3.0] - 2023-01-16
 

dataclass_array/array_dataclass.py

@@ -37,6 +37,9 @@
 import typing_extensions
 from typing_extensions import Annotated, Literal, TypeAlias  # pylint: disable=g-multiple-import
 
+if typing.TYPE_CHECKING:
+  import torch  # pytype: disable=import-error
+
 lazy = enp.lazy
 
 # TODO(pytype): Should use `dca.typing.DcT` but bound does not work across
@@ -477,8 +480,6 @@ def as_xnp(self: _DcT, xnp: enp.NpModule) -> _DcT:
     )
     return new_self
 
-  # ====== Internal ======
-
   # TODO(pytype): Remove hack. Currently, Python does not support typing
   # annotations for modules, by pytype auto-infer the correct type.
   # So this hack allow auto-completion
@@ -497,6 +498,33 @@ def xnp(self) -> enp.NpModule:
       """Returns the numpy module of the class (np, jnp, tnp)."""
       return self._xnp
 
+  # ====== Torch specific methods ======
+  # Could also add
+  # * x.detach
+  # * x.is_cuda
+  # * x.device
+  # * x.get_device
+
+  def to(self: _DcT, device, **kwargs) -> _DcT:
+    """Move the dataclass array to the device."""
+    if not lazy.is_torch_xnp(self.xnp):
+      raise ValueError('`.to` can only be called when `xnp == torch`')
+    return self.map_field(lambda f: f.to(device, **kwargs))
+
+  def cpu(self: _DcT, *args, **kwargs) -> _DcT:
+    """Move the dataclass array to the CPU device."""
+    if not lazy.is_torch_xnp(self.xnp):
+      raise ValueError('`.cpu` can only be called when `xnp == torch`')
+    return self.map_field(lambda f: f.cpu(*args, **kwargs))
+
+  def cuda(self: _DcT, *args, **kwargs) -> _DcT:
+    """Move the dataclass array to the CUDA device."""
+    if not lazy.is_torch_xnp(self.xnp):
+      raise ValueError('`.cuda` can only be called when `xnp == torch`')
+    return self.map_field(lambda f: f.cuda(*args, **kwargs))
+
+  # ====== Internal ======
+
   @epy.cached_property
   def _all_array_fields(self) -> dict[str, _ArrayField]:
     """All array fields, including `None` values."""

dataclass_array/array_dataclass_test.py

@@ -597,12 +597,19 @@ def test_infer_np(xnp: enp.NpModule):
         enp.lazy.torch.utils._pytree.tree_map,
     ],
 )
-def test_torch_tree_map(tree_map, dca_cls: DcaTest):
+def test_tree_map(tree_map, dca_cls: DcaTest):
   p = dca_cls.make(shape=(3,), xnp=np)
   p = tree_map(lambda x: x[None, ...], p)
   dca_cls.assert_val(p, (1, 3), xnp=np)
 
 
+def test_torch_device():
+  p = Nested.make(shape=(2,), xnp=enp.lazy.torch)
+  p = p.cpu()
+  p = p.to('cpu')
+  Nested.assert_val(p, (2,), xnp=enp.lazy.torch)
+
+
 @enp.testing.parametrize_xnp(
     restrict=[
         'jnp',