handle requires_grad in torchax (#8992)

Author: qihqi

torchax/test/test_functions.py

@@ -22,24 +22,23 @@ def forward(self, x: torch.Tensor):
 class TestTorchFunctions(parameterized.TestCase):
 
   def setUp(self):
-    self.env = torchax.tensor.Environment()
-    self.env.config.use_torch_native_for_cpu_tensor = False
+    torchax.enable_globally()
     torchax.enable_accuracy_mode()
+    self.env = torchax.default_env()
 
   @parameterized.named_parameters(
-      ('tensor_2d', lambda: torch.tensor([[0.1, 1.2], [2.2, 3.1], [4.9, 5.2]])),
-      ('tensor_1d', lambda: torch.tensor([0, 1],)),
-      ('tensor_scalar', lambda: torch.tensor(3.14159,)),
-      ('tensor_empty', lambda: torch.tensor([],)),
-      ('tensor_dtype', lambda: torch.tensor([[0.11111, 0.222222, 0.3333333]],
-                                            dtype=torch.float64)),
+      ('tensor_2d', [[0.1, 1.2], [2.2, 3.1], [4.9, 5.2]]),
+      ('tensor_1d', [0, 1]),
+      ('tensor_scalar', 3.14159),
+      ('tensor_empty', []),
+      ('tensor_dtype', [[0.11111, 0.222222, 0.3333333]], {'dtype': torch.float64})
   )
-  def test_tensor_constructor(self, func: Callable[[], torch.Tensor]):
-    expected = func()
+  def test_tensor_constructor(self, arg, kwargs=None):
+    kwargs = kwargs or {}
+    expected = torch.tensor(arg, **kwargs)
 
-    with self.env:
-      actual = func()
-      self.assertIsInstance(actual, torchax.tensor.Tensor)
+    actual = torch.tensor(arg, device='jax', **kwargs)
+    self.assertIsInstance(actual, torchax.tensor.Tensor)
 
     torch.testing.assert_close(torchax.tensor.j2t(actual._elem), expected)
 
@@ -90,8 +89,9 @@ def test_rms_norm(self):
       self.assertTrue(
         torch.allclose(res, torchax.tensor.j2t(res2.jax())))
 
-
-
+  def test_randn_requires_grad(self):
+    x = torch.randn((3, 3), requires_grad=True, device='jax')
+    self.assertEqual(x.requires_grad, True)
 
 
 if __name__ == '__main__':

torchax/torchax/tensor.py

@@ -411,12 +411,15 @@ def _handle_tensor_constructor(self, func, args, kwargs):
         with mode_utils.no_dispatch(), torch._C.DisableTorchFunction():
           return func(*args, **kwargs)
       with jax.default_device(jax_device):
+        requires_grad = kwargs.get('requires_grad', False)
         op = self._ops.get(func)
         if op is None and isinstance(func, torch._ops.OpOverload):
           op = self._ops.get(func.overloadpacket)
         res = op.func(*args, **kwargs)
         if isinstance(res, jax.Array):
           res = Tensor(res, self)
+        if requires_grad:
+          res.requires_grad = True
         return res
 
     def _torch_Tensor_to(self, args, kwargs):