enable tensor.copy_ with torch native cpu tensor as source

Author: yaochengji

torchax/test/test_functions.py

@@ -66,6 +66,13 @@ def test_flatten(self):
       a = a.flatten(0, 1)
       self.assertEqual(tuple(a.shape), (6, 4))
 
+  def test_copy_(self):
+    with self.env:
+      a = torch.zeros((2, 3), device="cpu")
+      b = torch.ones((2, 3))
+      b.copy_(a)
+      self.assertEqual(a, b.cpu())
+
   def test_rnn(self):
     model = SeqModel()
     x = torch.randn((2, 100, 20))

torchax/torchax/ops/jaten.py

@@ -122,6 +122,8 @@ def _aten_add(x, y, *, alpha=1):
 
 @op(torch.ops.aten.copy_, is_jax_function=False)
 def _aten_copy(x, y, memory_format=None):
+  if y.device.type == "cpu":
+    y = y.to(x.device)
   if x.ndim == 1 and y.ndim == 0:
     # case of torch.empty((1,)).copy_(tensor(N))
     # we need to return 0D tensor([N]) and not scalar tensor(N)

torchax/torchax/tensor.py

@@ -98,6 +98,13 @@ def shape(self):
   @property
   def ndim(self):
     return len(self._elem.shape)
+  
+  @property
+  def data(self):
+    return self
+
+  def copy_(self, other):
+    return torch.ops.aten.copy_(self, other)
 
   def flatten(self, start_dim=0, end_dim=-1):
     if end_dim == -1: