Add an option to not use dlpack.

DLpack introduces Jax arrays with wierd Layout.

Author: qihqi

benchmarks/util.py

@@ -72,10 +72,10 @@ def is_xla_device_available(devkind, use_xla2: bool = False):
 
 def move_to_device(item, device, torch_xla2: bool = False):
   if torch_xla2:
-    import torch_xla2
+    import torchax
     import jax
     move_to_device_func = lambda t: jax.device_put(
-        torch_xla2.tensor.t2j(t), device)
+        torchax.default_env().t2j_copy(t), device)
   else:
 
     def move_to_device_func(tensor: torch.Tensor) -> torch.Tensor:

torchax/examples/train_llama/utils.py

@@ -200,7 +200,8 @@ def _shard_fsdp_style(self, state_dict, sharding=None):
       sharding = self.x_sharding
 
     def move_one_tensor(x):
-      jval = torchax.tensor.t2j(x)
+      env = torchax.default_env()
+      jval = env.t2j_copy(x)
       return sharded_device_put(jval, sharding)
 
     if isinstance(state_dict, torch.Tensor):

torchax/test/gemma/test_gemma.py

@@ -71,7 +71,8 @@ def test_gemma(self):
     )
 
     weights, jax_func = torchax.extract_jax(model)
-    inputs_jax = pytree.tree_map_only(torch.Tensor, torchax.tensor.t2j, inputs)
+    env = torchax.default_env()
+    inputs_jax = env.t2j_copy(inputs)
 
     import jax
     print(jax.jit(jax_func)(weights, inputs_jax))

torchax/test/llama/test_llama.py

@@ -88,8 +88,8 @@ def make_cache(args, batch_size):
       m_prefill = torch.export.export(m, sample_input_prefill)
 
     weights, mj_prefill = torchax.export.exported_program_to_jax(m_prefill)
-    sample_inputs = pytree.tree_map_only(torch.Tensor, tensor.t2j,
-                                         sample_input_prefill)
+    env = torchax.default_env()
+    sample_inputs = env.t2j_copy(sample_input_prefill)
     print('Prefill', mj_prefill(weights, sample_inputs))
 
     sample_input_decode = (
@@ -103,8 +103,7 @@ def make_cache(args, batch_size):
     with torch.no_grad():
       m_decode = torch.export.export(m, sample_input_decode)
     weights, mj_decode = torchax.export.exported_program_to_jax(m_decode)
-    sample_inputs = pytree.tree_map_only(torch.Tensor, tensor.t2j,
-                                         sample_input_decode)
+    sample_inputs = env.t2j_copy(sample_input_decode)
     print('Decode', mj_decode(weights, sample_inputs))
 
 

torchax/test/moe/moe_test.py

@@ -48,7 +48,7 @@ def test_moe_layer(self):
     x_xla = env.to_xla(x)
     with jax.default_matmul_precision('float32'):
       res_xla = model_xla(x_xla)
-    res2 = torchax.tensor.j2t(res_xla._elem)
+    res2 = res_xla.to('cpu')
     print('max diff', torch.max((res - res2).abs()))
 
     self.assertTrue(torch.allclose(res2, res, atol=1e-2))

torchax/test/test_context.py

@@ -47,8 +47,7 @@ def test_same_manual_seed(self):
       y = torch.randn((3, 3))
       self.assertIsInstance(y, tensor.Tensor)
 
-    self.assertTrue(
-        torch.equal(torchax.tensor.j2t(x._elem), torchax.tensor.j2t(y._elem)))
+      self.assertTrue(torch.allclose(x, y))
 
   def test_different_manual_seed(self):
     with xla_env:
@@ -60,8 +59,7 @@ def test_different_manual_seed(self):
       y = torch.randn((3, 3))
       self.assertIsInstance(y, tensor.Tensor)
 
-    self.assertFalse(
-        torch.equal(torchax.tensor.j2t(x._elem), torchax.tensor.j2t(y._elem)))
+      self.assertFalse(torch.allclose(x, y))
 
   def test_jit_with_rng(self):
 
@@ -76,20 +74,14 @@ def random_op():
 
     # Result always expected to be the same for a jitted function because seeds
     # are baked in
-    torch.testing.assert_close(
-        torchax.tensor.j2t(random_jit()._elem),
-        torchax.tensor.j2t(random_jit()._elem),
-        atol=0,
-        rtol=0)
+    torch.testing.assert_close(random_jit(), random_jit(), atol=0, rtol=0)
 
   def test_generator_seed(self):
     with xla_env:
       x = torch.randn(2, 3, generator=torch.Generator().manual_seed(0))
       y = torch.randn(2, 3, generator=torch.Generator().manual_seed(0))
 
-    # Values will be different, but still check device, layout, dtype, etc
-    torch.testing.assert_close(
-        torchax.tensor.j2t(x._elem), torchax.tensor.j2t(y._elem))
+      torch.testing.assert_close(x, y)
 
   def test_buffer(self):
 

torchax/test/test_conv.py

@@ -55,24 +55,26 @@ def forward(self, x):
 class ConvTest(base_test_util.TestCase):
 
   def test_conv1(self):
+    env = torchax.default_env()
     m = CustomConv1()
     arg = torch.randn((20, 1, 50))
     res = m(arg)
 
     jax_weights, jax_func = torchax.extract_jax(m)
-    arg = torchax.tensor.t2j(arg)
+    arg = env.t2j_copy(arg)
     res2 = jax_func(jax_weights, (arg,))
-    res2_torch = torchax.tensor.j2t(res2)
+    res2_torch = env.j2t_copy(res2)
     self.assertTrue(torch.allclose(res, res2_torch))
 
   def test_conv2(self):
+    env = torchax.default_env()
     m = CustomConv2()
     arg = torch.randn((20, 4, 50, 100))
     res = m(arg)
     jax_weights, jax_func = torchax.extract_jax(m)
-    arg = torchax.tensor.t2j(arg)
+    arg = env.t2j_copy(arg)
     res2 = jax_func(jax_weights, (arg,))
-    res2_torch = torchax.tensor.j2t(res2)
+    res2_torch = env.j2t_copy(res2)
     self.assertTrue(torch.allclose(res, res2_torch, atol=1e-4, rtol=1e-4))
 
 

torchax/test/test_exports.py

@@ -43,12 +43,14 @@ def test_interpolate(self):
     model = Interpolate()
     ans = model(*arg)
 
+    env = torchax.default_env()
+
     with torch.no_grad():
       exported = torch.export.export(model, arg)
     weights, func = torchax.export.exported_program_to_jax(exported)
-    argj = tensor.t2j(arg[0])
+    argj = env.t2j_copy(arg[0])
     ans2 = jax.jit(func)(weights, (argj,))[0]
-    ans2 = tensor.j2t(ans2)
+    ans2 = env.j2t_copy(ans2)
     self.assertTrue(torch.allclose(ans, ans2, atol=1e-3))
 
     # Convert to StableHLO
@@ -67,11 +69,11 @@ def test_constant(self):
 
     with torch.no_grad():
       exported = torch.export.export(model, arg)
-
+    env = torchax.default_env()
     weights, func = torchax.export.exported_program_to_jax(exported)
-    argj = tensor.t2j(arg[0])
+    argj = env.t2j_copy(arg[0])
     ans2 = jax.jit(func)(weights, (argj,))[0]
-    ans2 = tensor.j2t(ans2)
+    ans2 = env.j2t_copy(ans2)
     self.assertTrue(torch.allclose(ans, ans2, atol=1e-5))
 
     # Convert to StableHLO

torchax/test/test_functions.py

@@ -40,7 +40,7 @@ def test_tensor_constructor(self, arg, kwargs=None):
     actual = torch.tensor(arg, device='jax', **kwargs)
     self.assertIsInstance(actual, torchax.tensor.Tensor)
 
-    torch.testing.assert_close(torchax.tensor.j2t(actual._elem), expected)
+    torch.testing.assert_close(actual.to('cpu'), expected)
 
   def test_dont_capture_conversion(self):
     t = torch.tensor([1, 2, 3])
@@ -86,7 +86,7 @@ def test_rms_norm(self):
       model.to('jax')
       x = x.to('jax')
       res2 = model(x)
-      self.assertTrue(torch.allclose(res, torchax.tensor.j2t(res2.jax())))
+      self.assertTrue(torch.allclose(res, res2.to('cpu')))
 
   def test_randn_requires_grad(self):
     x = torch.randn((3, 3), requires_grad=True, device='jax')

torchax/test/test_image.py

@@ -10,14 +10,6 @@
 import torchax.interop
 
 
-def to_xla_tensor(tensorstree):
-  return torchax.interop.torch_view(torchax.tensor.t2j(tensorstree))
-
-
-def to_torch_tensor(tensorstree):
-  return torchax.tensor.j2t(torchax.interop.jax_view(tensorstree))
-
-
 @partial(jax.jit, static_argnums=(1, 2, 3, 4))
 def upsample_jit(tensor, output_size: Tuple[int, int], align_corners: bool,
                  antialias: bool, method: str):
@@ -53,8 +45,9 @@ def test_resampling_combinations_bicubic(self, antialias, align_corners):
         align_corners=align_corners,
         antialias=antialias)
 
-    with torchax.default_env():
-      input_tensor_xla = to_xla_tensor(input_tensor)
+    env = torchax.default_env()
+    with env:
+      input_tensor_xla = env.to_xla(input_tensor)
       input_tensor_xla = torchax.interop.jax_view(input_tensor_xla)
       upsampled_tensor_xla = upsample_jit(
           input_tensor_xla,
@@ -63,7 +56,7 @@ def test_resampling_combinations_bicubic(self, antialias, align_corners):
           antialias=antialias,
           method=method)
 
-    upsampled_tensor_xla = to_torch_tensor(upsampled_tensor_xla)
+    upsampled_tensor_xla = env.j2t_copy(upsampled_tensor_xla)
     abs_err = torch.abs(upsampled_tensor - upsampled_tensor_xla)
 
     assert torch.allclose(