torchax 0.0.6__tar.gz → 0.0.7__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchax
-Version: 0.0.6
+Version: 0.0.7
 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 → torchax-0.0.7}/test/llama/llama_model.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
-# This file is copied from https://github.com/pytorch-labs/gpt-fast
+# This file is copied from https://github.com/meta-pytorch/gpt-fast
 # This is used for unit test purposes
 from dataclasses import dataclass
 import math

torchax-0.0.7/test/test_misc.py

@@ -0,0 +1,47 @@
+"""If you don't know which file a test should go, and don't want to make a new file
+for a small test. PUt it here
+"""
+import torch
+import unittest
+import torchax
+import jax
+import jax.numpy as jnp
+
+
+class MiscTest(unittest.TestCase):
+
+  def test_extract_jax_kwargs(self):
+
+    class M(torch.nn.Module):
+
+      def forward(self, a, b):
+        return torch.sin(a) + torch.cos(b)
+
+    weights, func = torchax.extract_jax(M())
+    res = func(
+        weights,
+        args=(),
+        kwargs={
+            'a': jnp.array([1, 2, 3]),
+            'b': jnp.array([3, 4, 5])
+        })
+    self.assertTrue(
+        jnp.allclose(
+            res,
+            jnp.sin(jnp.array([1, 2, 3])) + jnp.cos(jnp.array([3, 4, 5]))))
+
+  def test_to_device(self):
+    env = torchax.default_env()
+    env.config.debug_print_each_op = True
+    with env:
+      step1 = torch.ones(
+          100,
+          100,
+      )
+      step2 = torch.triu(step1, diagonal=1)
+      step3 = step2.to(dtype=torch.bool, device='jax')
+      self.assertEqual(step3.device.type, 'jax')
+
+
+if __name__ == '__main__':
+  unittest.main()

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

@@ -8,7 +8,7 @@ from torch.utils import _pytree as pytree
 from torchax import tensor
 from contextlib import contextmanager
 
-__version__ = "0.0.6"
+__version__ = "0.0.7"
 VERSION = __version__
 
 __all__ = [

{torchax-0.0.6 → torchax-0.0.7}/torchax/interop.py

@@ -237,6 +237,36 @@ def j2t_autograd(fn, call_jax=call_jax):
     the PyTorch autograd framework by saving the residuals into the context object.
     """
 
+  # NOTE(qihqi): This function cannot be inlined from the callsite
+  #  Becuase if it does, then it won't hit the compilation cache for
+  #  call_jax. Call jax uses functions' id as key.
+  # It is nested inside j2t_autograd to ensure it gets a unique ID for each
+  # wrapped pure function, preventing cache collisions between different pure modules.
+  def _jax_forward(fn, other, tree_def, tensors):
+    """JAX function to compute output and vjp function.
+  
+    primals should be a tuple (args, kwargs).
+    """
+    import jax
+    from jax.tree_util import tree_flatten, tree_unflatten
+
+    def fn_wrapper(*tensors):
+      # Reconstruct the original args and kwargs
+      flat_inputs = util.merge(tensors, other)
+      args, kwargs = tree_unflatten(tree_def, flat_inputs)
+      return fn(*args, **kwargs)
+
+    return jax.vjp(fn_wrapper, *tensors)
+
+  def _jax_backward(vjp_spec, saved_tensors, grad_out):
+    """JAX function to compute input gradients.
+  
+    Unflattening `saved_tensors` with `vjp_spec` should restore the original vjp function.
+    """
+    from jax.tree_util import tree_unflatten
+    fun_vjp = tree_unflatten(vjp_spec, saved_tensors)
+    return fun_vjp(grad_out)
+
   @wraps(fn)
   def inner(*args, **kwargs):
     from jax.tree_util import tree_flatten
@@ -290,36 +320,6 @@ def j2t_autograd(fn, call_jax=call_jax):
   return inner
 
 
-# NOTE(qihqi): This function cannot be inlined from the callsite
-#  Becuase if it does, then it won't hit the compilation cache for
-#  call_jax. Call jax uses functions' id as key.
-def _jax_forward(fn, other, tree_def, tensors):
-  """JAX function to compute output and vjp function.
-
-  primals should be a tuple (args, kwargs).
-  """
-  import jax
-  from jax.tree_util import tree_flatten, tree_unflatten
-
-  def fn_wrapper(*tensors):
-    # Reconstruct the original args and kwargs
-    flat_inputs = util.merge(tensors, other)
-    args, kwargs = tree_unflatten(tree_def, flat_inputs)
-    return fn(*args, **kwargs)
-
-  return jax.vjp(fn_wrapper, *tensors)
-
-
-def _jax_backward(vjp_spec, saved_tensors, grad_out):
-  """JAX function to compute input gradients.
-
-  Unflattening `saved_tensors` with `vjp_spec` should restore the original vjp function.
-  """
-  from jax.tree_util import tree_unflatten
-  fun_vjp = tree_unflatten(vjp_spec, saved_tensors)
-  return fun_vjp(grad_out)
-
-
 fori_loop = torch_view(jax.lax.fori_loop)
 
 

{torchax-0.0.6 → torchax-0.0.7}/torchax/ops/mappings.py

@@ -6,6 +6,14 @@ import torch.func
 import torch.utils.dlpack as torchdl
 import torch.utils._mode_utils as mode_utils
 
+NUMPY_UNSUPPORTED_DTYPES = {
+    torch.bfloat16: jnp.bfloat16,
+    torch.float8_e4m3fn: jnp.float8_e4m3fn,
+    torch.float8_e4m3fnuz: jnp.float8_e4m3fnuz,
+    torch.float8_e5m2: jnp.float8_e5m2,
+    torch.float8_e5m2fnuz: jnp.float8_e5m2fnuz,
+}
+
 
 def t2j(t, use_dlpack=True):
   is_bool = False
@@ -28,14 +36,14 @@ def t2j(t, use_dlpack=True):
   if res is None:
     # https://github.com/google/jax/issues/7657
     # https://github.com/google/jax/issues/17784
-    if t.dtype == torch.bfloat16:
+    if t.dtype in NUMPY_UNSUPPORTED_DTYPES:
       nparray = (t.cpu().detach().to(torch.float32).numpy()
-                )  # numpy don't support bfloat16
+                )  # handle dtypes not supported by numpy
     else:
       nparray = t.cpu().detach().numpy()
     res = jnp.asarray(nparray)
-    if t.dtype == torch.bfloat16:
-      res = res.astype(jnp.bfloat16)
+    if t.dtype in NUMPY_UNSUPPORTED_DTYPES:
+      res = res.astype(NUMPY_UNSUPPORTED_DTYPES[t.dtype])
 
   if is_bool:
     res = res.astype(jnp.bool_)

torchax-0.0.7/torchax/ops/type_casting.py

@@ -0,0 +1,47 @@
+
+from torchax import config
+from torchax.ops import mappings
+import jax.numpy as jnp
+import torch
+
+def maybe_cast(result, torch_op):
+  """Casts the result to the torch op's return dtype if the config is set."""
+  if not config.DEFAULTS.internal_respect_torch_return_dtypes:
+    return result
+
+  if not hasattr(torch_op, '_schema'):
+    return result
+
+  schema = torch_op._schema
+  if not schema.returns:
+    return result
+
+  # TODO: Handle multiple return values
+  if len(schema.returns) > 1:
+    return result
+
+  return_type = schema.returns[0].type
+  if str(return_type) == 'Tensor':
+    # This is not quite right, we need to get the dtype of the tensor
+    # For now, let's assume we can get it from the first input argument
+    if not schema.arguments:
+      return result
+    
+    input_type = schema.arguments[0].type
+    if str(input_type) != 'Tensor':
+      return result
+      
+    # This is a hack, we need a better way to determine the return dtype
+    # For now, let's assume the return type is the same as the first input
+    # This is not always true, e.g. for comparison ops.
+    return result
+
+  try:
+    torch_dtype = getattr(torch, str(return_type))
+    jax_dtype = mappings.t2j_dtype(torch_dtype)
+    if isinstance(result, jnp.ndarray):
+      return result.astype(jax_dtype)
+    else:
+      return jax_dtype(result)
+  except (AttributeError, TypeError):
+    return result

{torchax-0.0.6 → torchax-0.0.7}/torchax/tensor.py

@@ -469,12 +469,12 @@ class Environment(contextlib.ContextDecorator):
         arr = self.t2j_copy(the_tensor)
         res = Tensor(arr, self, the_tensor.requires_grad)
 
-    if new_dtype is not None and new_dtype != the_tensor.dtype:
-      if isinstance(the_tensor, Tensor):
+    if new_dtype is not None and new_dtype != res.dtype:
+      if isinstance(res, Tensor):
         res = res.apply_jax(jnp.astype, mappings.t2j_dtype(new_dtype))
       else:
         with mode_utils.no_dispatch(), torch._C.DisableTorchFunction():
-          return the_tensor.to(device=new_device, dtype=new_dtype)
+          return res.to(device=new_device, dtype=new_dtype)
     return res
 
   def get_and_rotate_prng_key(self,

torchax-0.0.6/test/test_misc.py

@@ -1,22 +0,0 @@
-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()