Bug fixes (#9554)

Author: qihqi

torchax/test/test_mutations.py

@@ -8,28 +8,44 @@ class TestMutations(TestCase):
 
   def setUp(self):
     self.env = torchax.tensor.Environment()
+    self.env.config.debug_print_each_op = True
 
   def test_add(self):
     with self.env:
-      x = torch.tensor([1, 2, 3], dtype=torch.int32)
-      y = torch.tensor([4, 5, 6], dtype=torch.int32)
+      x = torch.tensor([1, 2, 3], device='jax', dtype=torch.int32)
+      y = torch.tensor([4, 5, 6], device='jax', dtype=torch.int32)
       x.add_(y)
-      self.assertEqual(x, torch.tensor([5, 7, 9], dtype=torch.int32))
+      torch.testing.assert_close(x.cpu(),
+                                 torch.tensor([5, 7, 9], dtype=torch.int32))
 
   def test_sub(self):
     with self.env:
-      x = torch.tensor([1, 2, 3], dtype=torch.int32)
-      y = torch.tensor([4, 5, 6], dtype=torch.int32)
+      x = torch.tensor([1, 2, 3], device='jax', dtype=torch.int32)
+      y = torch.tensor([4, 5, 6], device='jax', dtype=torch.int32)
       x.sub_(y)
-      self.assertEqual(x, torch.tensor([-3, -3, -3], dtype=torch.int32))
+      torch.testing.assert_close(x.cpu(),
+                                 torch.tensor([-3, -3, -3], dtype=torch.int32))
 
   def test_mul(self):
     with self.env:
-      x = torch.tensor([1, 2, 3], dtype=torch.int32)
-      y = torch.tensor([4, 5, 6], dtype=torch.int32)
+      x = torch.tensor([1, 2, 3], device='jax', dtype=torch.int32)
+      y = torch.tensor([4, 5, 6], device='jax', dtype=torch.int32)
 
       x.mul_(y)
-      self.assertEqual(x, torch.tensor([4, 10, 18], dtype=torch.int32))
+      torch.testing.assert_close(x.cpu(),
+                                 torch.tensor([4, 10, 18], dtype=torch.int32))
+
+  def test_index_copy(self):
+    with self.env:
+      x = torch.zeros(5, 3, device='jax')
+      t = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+                       device='jax',
+                       dtype=torch.float)
+      index = torch.tensor([0, 4, 2], device='jax')
+      x.index_copy_(0, index, t)
+      expected = torch.tensor([[1., 2., 3.], [0., 0., 0.], [7., 8., 9.],
+                               [0., 0., 0.], [4., 5., 6.]])
+      torch.testing.assert_close(x.cpu(), expected)
 
 
 if __name__ == '__main__':

torchax/torchax/ops/jaten.py

@@ -736,7 +736,6 @@ def _aten_empty_strided(sizes, stride, dtype=None, **kwargs):
   return jnp.empty(sizes, dtype=dtype)
 
 
-@op(torch.ops.aten.index_put_)
 @op(torch.ops.aten.index_put)
 def _aten_index_put(self, indexes, values, accumulate=False):
   indexes = [slice(None, None, None) if i is None else i for i in indexes]
@@ -5618,6 +5617,8 @@ def _aten__assert_tensor_metadata(*args, **kwargs):
         op_base.InplaceOp(torch.ops.aten.floor_divide),
     torch.ops.aten.remainder_:
         op_base.InplaceOp(torch.ops.aten.remainder),
+    torch.ops.aten.index_put_:
+        op_base.InplaceOp(torch.ops.aten.index_put),
 }
 
 # Note: tuple comparisons work intuitively, e.g. `_jax_version >= (0, 4, 32)`.

torchax/torchax/tensor.py

@@ -70,9 +70,6 @@ def __str__(self):
 
   __repr__ = __str__
 
-  def __jax_array__(self):
-    return self._elem
-
   @property
   def shape(self):
     return torch.Size(self._elem.shape)
@@ -494,6 +491,8 @@ def _handle_tensor_constructor(self, func, args, kwargs):
       op = self._get_op_or_decomp(func)
       if op.needs_env:
         kwargs['env'] = self
+      if op.is_jax_function:
+        (args, kwargs) = self.t2j_iso((args, kwargs))
       res = op.func(*args, **kwargs)
       if isinstance(res, jax.Array):
         res = Tensor(res, self, requires_grad)