torchax 0.0.5__tar.gz → 0.0.6__tar.gz

{torchax-0.0.5 → torchax-0.0.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchax
-Version: 0.0.5
+Version: 0.0.6
 Summary: torchax is a library for running Jax and PyTorch together
 Project-URL: Homepage, https://github.com/pytorch/xla/tree/master/torchax
 Author-email: Han Qi <qihan.dev@gmail.com>, Pytorch/XLA team <pytorchxla-dev@google.com>

torchax-0.0.6/docs/api_iterations.md

@@ -0,0 +1,22 @@
+
+## always create a new environment, use it, discard it?
+
+```
+env = torchax.env()
+with env.with_prng_key(): // or extra inputs
+  do stuff
+
+# discard env
+
+with env.output_shape
+with env.manual axis ...
+
+env.call_torch_func(f, args, kwargs)
+
+functions should take in env
+functions in torch will get env from threadlocal property
+```
+env.call_stateless_torch_func()?
+
+tx = torchax.initialize(...)
+```

torchax-0.0.6/test/BUILD

@@ -0,0 +1,31 @@
+# TODO(hanq): describe this package.
+
+load(
+    "//third_party/py/torch/google/bazel_rules/rules_python/python:defs.bzl",
+    "py_library",
+    "py_test",
+)
+
+package(
+    default_applicable_licenses = ["//devtools/compliance/licenses:no_external_contributions"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "test_base",
+    srcs = ["test_base.py"],
+    deps = [
+        "//testing/pybase",
+    ],
+)
+
+py_test(
+    name = "test_core_aten_ops",
+    srcs = ["test_core_aten_ops.py"],
+    deps = [
+        ":test_base",
+        "//third_party/py/absl:app",
+        "//third_party/py/torch/google/_torx",
+    ],
+)

torchax-0.0.6/test/test_base.py

@@ -0,0 +1,55 @@
+import unittest
+import torch
+from torch.utils import _pytree as pytree
+
+from torchax import tensor
+
+TestCase = unittest.TestCase
+main = unittest.main
+
+
+def diff_output(testcase, output1, output2, rtol, atol, equal_nan=True):
+  if isinstance(output1, torch.Tensor):
+    testcase.assertIsInstance(output2, torch.Tensor)
+    output2_cpu = output2.detach().cpu()
+    if output2_cpu.dtype != output1.dtype:
+      output2_cpu = output2_cpu.to(output1.dtype)
+    testcase.assertTrue(
+        torch.allclose(
+            output1, output2_cpu, atol=atol, rtol=rtol, equal_nan=equal_nan))
+  elif isinstance(output1, (tuple, list)):
+    testcase.assertIsInstance(output2, (tuple, list))
+    testcase.assertEqual(len(output1), len(output2))
+    for o1, o2 in zip(output1, output2):
+      diff_output(testcase, o1, o2, rtol, atol)
+  else:
+    testcase.assertEqual(output1, output2)
+
+
+def run_function_and_compare(testcase,
+                             func,
+                             args,
+                             kwargs,
+                             atol=1e-3,
+                             rtol=1e-5,
+                             equal_nan=True,
+                             ignore_indices=False):
+  with testcase.subTest("torch_eval"):
+    res = func(*args, **kwargs)
+    with testcase.subTest("torchax_eval"):
+      args2, kwargs2 = pytree.tree_map_only(torch.Tensor, tensor.move_to_device,
+                                            (args, kwargs))
+      res2 = func(*args2, **kwargs2)
+      res2 = pytree.tree_map_only(tensor.Tensor, lambda t: t.torch(), res2)
+      with testcase.subTest("torchax_diff:" + str(atol)):
+        if ignore_indices and isinstance(res, tuple) and len(res) == 2:
+          diff_output(
+              testcase,
+              res[0],
+              res2[0],
+              atol=atol,
+              rtol=rtol,
+              equal_nan=equal_nan)
+        else:
+          diff_output(
+              testcase, res, res2, atol=atol, rtol=rtol, equal_nan=equal_nan)

{torchax-0.0.5 → torchax-0.0.6}/test/test_context.py

@@ -10,16 +10,9 @@ xla_env = torchax.default_env()
 
 class TestContext(unittest.TestCase):
 
-  def setUp(self):
-    self.old_var = xla_env.config.use_torch_native_for_cpu_tensor
-    xla_env.config.use_torch_native_for_cpu_tensor = False
-
-  def tearDown(self):
-    xla_env.config.use_torch_native_for_cpu_tensor = self.old_var
-
   def test_mode_context_manager(self):
     with xla_env:
-      x = torch.full((3, 3), -1)
+      x = torch.full((3, 3), -1, device='jax')
       self.assertIsInstance(x, tensor.Tensor)
       y = x.abs()
       self.assertIsInstance(y, tensor.Tensor)
@@ -27,7 +20,7 @@ class TestContext(unittest.TestCase):
   @staticmethod
   @xla_env
   def _test_mode_decorator():
-    x = torch.full((3, 3), -1)
+    x = torch.full((3, 3), -1).to('jax')
     y = x.abs()
 
     return x, y
@@ -40,11 +33,11 @@ class TestContext(unittest.TestCase):
   def test_same_manual_seed(self):
     with xla_env:
       xla_env.manual_seed(1234)
-      x = torch.randn((3, 3))
+      x = torch.randn((3, 3), device='jax')
       self.assertIsInstance(x, tensor.Tensor)
 
       xla_env.manual_seed(1234)
-      y = torch.randn((3, 3))
+      y = torch.randn((3, 3), device='jax')
       self.assertIsInstance(y, tensor.Tensor)
 
       self.assertTrue(torch.allclose(x, y))
@@ -52,11 +45,11 @@ class TestContext(unittest.TestCase):
   def test_different_manual_seed(self):
     with xla_env:
       xla_env.manual_seed(1234)
-      x = torch.randn((3, 3))
+      x = torch.randn((3, 3), device='jax')
       self.assertIsInstance(x, tensor.Tensor)
 
       xla_env.manual_seed(12345)
-      y = torch.randn((3, 3))
+      y = torch.randn((3, 3), device='jax')
       self.assertIsInstance(y, tensor.Tensor)
 
       self.assertFalse(torch.allclose(x, y))
@@ -66,21 +59,24 @@ class TestContext(unittest.TestCase):
     with xla_env:
 
       def random_op():
-        x = torch.randn(3, 3)
-        y = torch.randn(3, 3)
+        x = torch.randn(3, 3, device='jax')
+        y = torch.randn(3, 3, device='jax')
         return x @ y
 
       random_jit = torchax.interop.jax_jit(random_op)
       self.assertIsInstance(random_jit(), tensor.Tensor)
 
       # If we run the JIT twice, the random values should be different.
-      with self.assertRaises(AssertionError):
-        torch.testing.assert_close(random_jit(), random_jit(), atol=0, rtol=0)
+      # TODO(qihqi): think about API for passing down seed
+      #  with self.assertRaises(AssertionError):
+      #    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))
+      x = torch.randn(
+          2, 3, generator=torch.Generator().manual_seed(0), device='jax')
+      y = torch.randn(
+          2, 3, generator=torch.Generator().manual_seed(0), device='jax')
 
       # Values will be the same given the same seed.
       torch.testing.assert_close(x, y)
@@ -97,7 +93,7 @@ class TestContext(unittest.TestCase):
 
     # Test context manager.
     with xla_env:
-      m = M()
+      m = M().to('jax')
       self.assertIsInstance(m.c, tensor.Tensor)
       self.assertIsInstance(m.c2, tensor.Tensor)
     # Test `to_xla`.

{torchax-0.0.5 → torchax-0.0.6}/test/test_core_aten_ops.py

@@ -90,11 +90,6 @@ class TestCoreAtenOps(unittest.TestCase):
     super().setUp()
     torch.manual_seed(0)
     self.env = tensor.Environment()
-    self.old_var = self.env.config.use_torch_native_for_cpu_tensor
-    self.env.config.use_torch_native_for_cpu_tensor = False
-
-  def tearDown(self):
-    self.env.config.use_torch_native_for_cpu_tensor = self.old_var
 
   def test_aten_abs_0(self):
     args = (torch.randn((10, 10)).to(torch.float32),)

{torchax-0.0.5 → torchax-0.0.6}/test/test_flax.py

@@ -81,7 +81,7 @@ class FlaxTest(unittest.TestCase):
         return res
 
     with env:
-      nn_module = Parent()
+      nn_module = Parent().to('jax')
 
       @jax_jit
       def jitted(weights, args):

{torchax-0.0.5 → torchax-0.0.6}/test/test_functions.py

@@ -88,8 +88,15 @@ class TestTorchFunctions(parameterized.TestCase):
       res2 = model(x)
       self.assertTrue(torch.allclose(res, res2.to('cpu')))
 
-  def test_randn_requires_grad(self):
-    x = torch.randn((3, 3), requires_grad=True, device='jax')
+  @parameterized.named_parameters(
+      ('ones', torch.ones, ((2, 2),)), ('zeros', torch.zeros, ((2, 2),)),
+      ('empty', torch.empty,
+       ((2, 2),)), ('empty_strided', torch.empty_strided,
+                    ((2, 2), (2, 1))), ('tensor', torch.tensor, ([2.0, 2.0],)),
+      ('eye', torch.eye, (2,)), ('randn', torch.randn, ((2, 2),)),
+      ('rand', torch.rand, ((2, 2),)), ('full', torch.full, ((2, 2), 0)))
+  def test_requires_grad(self, func, args):
+    x = func(*args, requires_grad=True, device='jax')
     self.assertEqual(x.requires_grad, True)
 
 

{torchax-0.0.5 → torchax-0.0.6}/test/test_interop.py

@@ -2,9 +2,10 @@ import functools
 import torch
 import unittest
 import torchax
-from torchax import interop, jax_device
+from torchax import interop
 import torchax
 import jax
+import jax.numpy as jnp
 
 
 def is_tpu_available():
@@ -142,7 +143,7 @@ class InteropTest(unittest.TestCase):
       self.assertEqual(e.jax_device.platform, "cpu")
       self.assertEqual(e.device.type, "jax")
 
-    with jax_device("cpu"):
+    with jax.default_device(jax.devices("cpu")[0]):
       # move torch.tensor to torchax.tensor CPU
       b = a.to("jax")
       self.assertEqual(b.jax_device.platform, "cpu")
@@ -150,26 +151,21 @@ class InteropTest(unittest.TestCase):
 
     if is_tpu_available():
       # move torch.tensor to torchax.tensor TPU
-      with jax_device("tpu"):
+      with jax.default_device(jax.local_devices("tpu")[0]):
         c = a.to("jax")
         self.assertEqual(c.jax_device.platform, "tpu")
         self.assertEqual(c.device.type, "jax")
 
-      # move torchax.tensor on CPU to TPU
-      with jax_device("tpu"):
-        self.assertEqual(b.jax_device.platform, "cpu")
-        self.assertEqual(c.device.type, "jax")
-        c = b.to("jax")
-        self.assertEqual(c.jax_device.platform, "tpu")
-        self.assertEqual(c.device.type, "jax")
-
-      # move torchax.tensor on TPU to CPU
-      with jax_device("cpu"):
-        self.assertEqual(c.jax_device.platform, "tpu")
-        self.assertEqual(c.device.type, "jax")
-        d = c.to("jax")
-        self.assertEqual(d.jax_device.platform, "cpu")
-        self.assertEqual(d.device.type, "jax")
+  def test_torch_jax_view_dtype(self):
+    dtype = torch.float32
+    self.assertEqual(interop.jax_view(dtype), jnp.float32.dtype)
+    self.assertEqual(interop.torch_view(interop.jax_view(dtype)), dtype)
+    dtype = torch.bfloat16
+    self.assertEqual(interop.jax_view(dtype), jnp.bfloat16.dtype)
+    self.assertEqual(interop.torch_view(interop.jax_view(dtype)), dtype)
+    dtype = torch.int32
+    self.assertEqual(interop.jax_view(dtype), jnp.int32.dtype)
+    self.assertEqual(interop.torch_view(interop.jax_view(dtype)), dtype)
 
 
 if __name__ == '__main__':

{torchax-0.0.5 → torchax-0.0.6}/test/test_jittable_module.py

@@ -34,6 +34,25 @@ class JittableModuleTest(unittest.TestCase):
     assert isinstance(JittableMoreAwesomeModel, EvenMoreAwesomeModel)
     assert not isinstance(JittableMoreAwesomeModel, MyAwesomeModel)
 
+  def test_functional_call_callable(self):
+
+    def outer_function(model, x):
+      return x + 1
+
+    model = MyAwesomeModel()
+    jittable_module = interop.JittableModule(model)
+
+    # Check if the jittable module can be called like a function
+    input_tensor = torch.randn(1, 3, 224, 224)
+    expected_output = input_tensor + 1
+
+    output = jittable_module.functional_call(outer_function,
+                                             jittable_module.params,
+                                             jittable_module.buffers,
+                                             input_tensor)
+
+    assert torch.equal(output, expected_output)
+
 
 if __name__ == '__main__':
   unittest.main()

{torchax-0.0.5 → torchax-0.0.6}/test/test_libraries.py

@@ -54,7 +54,6 @@ class LibraryTest(unittest.TestCase):
 
   def setUp(self):
     torch.manual_seed(0)
-    torchax.default_env().config.use_torch_native_for_cpu_tensor = False
 
   def test_basic_sdpa_library(self):
 

torchax-0.0.6/test/test_misc.py

@@ -0,0 +1,22 @@
+import unittest
+import torch
+import torchax
+
+
+class MiscTest(unittest.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    torchax.enable_globally()
+
+  def test_mixed_tensor_math_with_scalar(self):
+    a = torch.tensor(2)
+    b = torch.ones((2, 2), device='jax')
+    c = a * b
+    self.assertTrue(
+        torch.allclose(c.cpu(),
+                       torch.tensor([[2, 2], [2, 2]], dtype=torch.float32)))
+
+
+if __name__ == '__main__':
+  unittest.main()

torchax-0.0.6/test/test_mutations.py

@@ -0,0 +1,52 @@
+import unittest
+import torchax
+import torch
+from torch.testing._internal.common_utils import TestCase
+
+
+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], device='jax', dtype=torch.int32)
+      y = torch.tensor([4, 5, 6], device='jax', dtype=torch.int32)
+      x.add_(y)
+      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], device='jax', dtype=torch.int32)
+      y = torch.tensor([4, 5, 6], device='jax', dtype=torch.int32)
+      x.sub_(y)
+      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], device='jax', dtype=torch.int32)
+      y = torch.tensor([4, 5, 6], device='jax', dtype=torch.int32)
+
+      x.mul_(y)
+      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__':
+  unittest.main()

{torchax-0.0.5 → torchax-0.0.6}/test/test_ops.py

@@ -140,6 +140,8 @@ def run_export_and_compare(testcase,
     with testcase.subTest("torchax_eval"):
       input2, args2, kwargs2 = testcase.env.to_xla(
           (sample_input.input, sample_input.args, sample_input.kwargs))
+      if 'device' in kwargs2:
+        kwargs2['device'] = 'jax'
       with testcase.env:
         res2 = func(input2, *args2, **kwargs2)
         res2 = pytree.tree_map_only(tensor.Tensor, lambda t: t.torch(), res2)
@@ -186,13 +188,8 @@ class TestOpInfo(TestCase):
     self.env = torchax.default_env()
     torchax.enable_accuracy_mode()
     #self.env.config.debug_accuracy_for_each_op = True
-    self.env.config.debug_print_each_op = True
+    self.env.config.debug_print_each_op = False
     torch.manual_seed(0)
-    self.old_var = self.env.config.use_torch_native_for_cpu_tensor
-    self.env.config.use_torch_native_for_cpu_tensor = False
-
-  def tearDown(self):
-    self.env.config.use_torch_native_for_cpu_tensor = self.old_var
 
   # Replaces all values in the input torch_tensor that are less than the given threshold
   # with the threshold value itself.

{torchax-0.0.5 → torchax-0.0.6}/test/test_unbounded_dynamism.py

@@ -53,13 +53,8 @@ def wrap_func_as_nn_module(f):
 class UnboundedDynamismExportTest(unittest.TestCase):
 
   def setUp(self):
-    self.env = torchax.default_env()
-    self.env.config.use_torch_native_for_cpu_tensor = False
     torchax.enable_accuracy_mode()
 
-  def tearDown(self):
-    self.env.config.use_torch_native_for_cpu_tensor = True
-
   def test_add(self):
     args = (torch.rand((10, 197, 768)), torch.rand((10, 197, 768)))
     dynamic_shapes = (({0: Dim("dim")}, {0: Dim("dim")}),)

{torchax-0.0.5 → torchax-0.0.6}/torchax/__init__.py

@@ -6,10 +6,9 @@ import os
 import torch
 from torch.utils import _pytree as pytree
 from torchax import tensor
-from torchax import distributed  # noqa: F401
 from contextlib import contextmanager
 
-__version__ = "0.0.5"
+__version__ = "0.0.6"
 VERSION = __version__
 
 __all__ = [
@@ -50,10 +49,11 @@ def extract_jax(mod: torch.nn.Module, env=None):
   states = env.t2j_copy(states)
 
   #@jax.jit
-  def jax_func(states, inputs):
-    (states, inputs) = env.j2t_iso((states, inputs))
+  def jax_func(states, args, kwargs=None):
+    (states, args, kwargs) = env.j2t_iso((states, args, kwargs))
     with env:
-      res = torch.func.functional_call(mod, states, inputs, tie_weights=False)
+      res = torch.func.functional_call(
+          mod, states, args, kwargs, tie_weights=False)
     return env.t2j_iso(res)
 
   return states, jax_func
@@ -81,11 +81,6 @@ def disable_temporarily():
 
 torch.utils.rename_privateuse1_backend('jax')
 unsupported_dtype = [torch.quint8]
-torch.utils.generate_methods_for_privateuse1_backend(
-    for_tensor=True,
-    for_module=True,
-    for_storage=True,
-    unsupported_dtype=unsupported_dtype)
 
 import jax
 import torchax.device_module
@@ -129,34 +124,3 @@ def compile(fn, options: Optional[CompileOptions] = None):
     raise RuntimeError('dynamo mode is not supported yet')
   elif options.mode == 'export':
     raise RuntimeError('export mode is not supported yet')
-
-
-@contextmanager
-def jax_device(target_device: str, env: tensor.Environment | None = None):
-  """
-  to("jax") cannot differentiate the device/platform (cpu vs tpu). 
-  Use this context manager to control jax array's storage device
-  
-  Examples:
-
-  a = torch.ones(3, 3)
-
-  with jax_device("cpu"):
-    b = a.to("jax") 
-  
-  with jax_device("tpu"):
-    c = a.to("jax") 
-  
-  with jax_device("tpu"):
-    c = b.to("jax") 
-
-  """
-  if env is None:
-    env = default_env()
-
-  prev_target_device = env.target_device
-  try:
-    env.target_device = target_device
-    yield env
-  finally:
-    env.target_device = prev_target_device

{torchax-0.0.5 → torchax-0.0.6}/torchax/amp.py

@@ -61,9 +61,8 @@ def autocast(device, dtype=torch.bfloat16, env=None):
   if env is None:
     import torchax
     env = torchax.default_env()
-  env.autocast_dtype, old = dtype, env.autocast_dtype
-  yield
-  env.autocast_dtype = old
+  with env.override_property(autocast_dtype=dtype):
+    yield
 
 
 # https://github.com/pytorch/pytorch/blob/05faba40287cf7d8734da96cb2e904f39710bf29/aten/src/ATen/autocast_mode.cpp#L327

{torchax-0.0.5 → torchax-0.0.6}/torchax/config.py

@@ -10,6 +10,11 @@ class Configuration:
 
   use_int32_for_index: bool = False
 
+  # normally, math between CPU torch.Tensor with torchax.Tensor is not
+  # allowed. However, if that torch.Tensor happens to be scalar, then we
+  # can use scalar * tensor math to handle it
+  allow_mixed_math_with_scalar_tensor: bool = True
+
   # If true, we will convert Views into torchax.Tensors eagerly
   force_materialize_views: bool = False
 
@@ -22,5 +27,4 @@ class Configuration:
 
   # device
   treat_cuda_as_jax_device: bool = True
-  use_torch_native_for_cpu_tensor: bool = True
   internal_respect_torch_return_dtypes: bool = False

torchax-0.0.6/torchax/configuration.py

@@ -0,0 +1,30 @@
+import dataclasses
+
+
+@dataclasses.dataclass
+class Configuration:
+  debug_print_each_op: bool = False
+  debug_accuracy_for_each_op: bool = False
+  debug_mixed_tensor: bool = False
+  debug_print_each_op_operands: bool = False
+
+  use_int32_for_index: bool = False
+
+  # normally, math between CPU torch.Tensor with torchax.Tensor is not
+  # allowed. However, if that torch.Tensor happens to be scalar, then we
+  # can use scalar * tensor math to handle it
+  allow_mixed_math_with_scalar_tensor: bool = True
+
+  # If true, we will convert Views into torchax.Tensors eagerly
+  force_materialize_views: bool = False
+
+  # Use DLPack for converting jax.Arrays <-> and torch.Tensor
+  use_dlpack_for_data_conversion: bool = False
+
+  # Flash attention
+  use_tpu_flash_attention: bool = False
+  shmap_flash_attention: bool = False
+
+  # device
+  treat_cuda_as_jax_device: bool = True
+  internal_respect_torch_return_dtypes: bool = False

{torchax-0.0.5 → torchax-0.0.6}/torchax/device_module.py

@@ -1,3 +1,6 @@
+import torch
+
+
 def _is_in_bad_fork():
   return False
 
@@ -24,3 +27,7 @@ def is_available():
 
 def current_device():
   return 0
+
+
+def get_amp_supported_dtype():
+  return [torch.float16, torch.bfloat16]

torchax-0.0.6/torchax/environment.py

@@ -0,0 +1 @@
+