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/tensor.py

@@ -95,6 +95,15 @@ def shape(self):
   @property
   def ndim(self):
     return len(self._elem.shape)
+  
+  @property
+  def data(self):
+    return self
+
+  def copy_(self, other):
+    if other.device.type == "cpu":
+      other = other.to(self.device)
+    return torch.ops.aten.copy_(self, other)
 
   def flatten(self, start_dim=0, end_dim=-1):
     if end_dim == -1: