torchax 0.0.4__py3-none-any.whl

torchax/tensor.py

@@ -0,0 +1,557 @@
+import logging
+import sys
+import contextlib
+from typing import Optional, Any
+import jax
+import jax.numpy as jnp
+import numpy
+import torch
+import torch.distributed._functional_collectives
+import torch.func
+import torch.utils._mode_utils as mode_utils
+import torch.utils._python_dispatch as torch_dispatch
+import torch.utils._pytree as torch_pytree
+
+from torchax import config
+from torchax.ops import mappings, ops_registry
+
+logger = logging.getLogger(__name__)
+
+
+class OperatorNotFound(Exception):
+  pass
+
+
+def wrap(jaxarray):
+  return torch_pytree.tree_map_only(jnp.ndarray, Tensor, jaxarray)
+
+
+def unwrap(torchtensors):
+  return torch_pytree.tree_map_only(Tensor, lambda x: x._elem, torchtensors)
+
+
+def t2j(t):
+  if isinstance(t, Tensor):
+    return t._elem
+  return mappings.t2j(t)
+
+
+def j2t(x):
+  return mappings.j2t(x)
+
+
+def t2j_dtype(dtype):
+  return mappings.t2j_dtype(dtype)
+
+
+def j2t_dtype(dtype):
+  return mappings.j2t_dtype(dtype)
+
+
+@contextlib.contextmanager
+def log_nested(env, message):
+  if env.config.debug_print_each_op:
+    print((' ' * log_nested.level) + message, file=sys.stderr)
+  log_nested.level += 1
+  yield
+  log_nested.level -= 1
+
+log_nested.level = 0
+
+
+class Tensor(torch.Tensor):
+
+  @staticmethod
+  def __new__(cls, elem, env):
+    dtype = j2t_dtype(elem.dtype)
+    shape = list(elem.shape)
+    for i, s in enumerate(shape):
+      if not isinstance(s, int):
+        shape[i] = 1
+    if dtype is None:
+      dtype = torch.float32
+    return torch.Tensor._make_wrapper_subclass(
+        cls,
+        shape,
+        dtype=dtype,
+        device='meta',
+        requires_grad=False,
+    )
+
+  def __init__(self, elem: jax.Array, env: 'Environment'):
+    super().__init__()
+    self._elem = elem
+    self._env = env
+
+  def __str__(self):
+    return "Tensor({} {})".format(str(type(self._elem)), str(self._elem))
+
+  __repr__ = __str__
+
+  def __jax_array__(self):
+    return self._elem
+
+  @property
+  def shape(self):
+    return self._elem.shape
+
+  @property
+  def ndim(self):
+    return len(self._elem.shape)
+
+  def flatten(self, start_dim=0, end_dim=-1):
+    if end_dim == -1:
+      end_dim = self.ndim
+    new_shape = (
+        self._elem.shape[:start_dim] + (-1,) + self._elem.shape[end_dim + 1:])
+    new_elem = jnp.reshape(self._elem, new_shape)
+    return Tensor(new_elem, self._env)
+    # return torch.reshape(self, new_shape)
+
+  def __setitem__(self, key, val):
+    key, val = self._env.t2j_iso((key, val))
+    self._elem = self._elem.at[key].set(val)
+
+  def type_as(self, other):
+    self._elem = self._elem.astype(other._elem.dtype)
+    return self
+
+  __torch_function__ = torch._C._disabled_torch_function_impl
+
+  @classmethod
+  def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
+    env = None
+    for arg in torch_pytree.arg_tree_leaves(*args, **kwargs):
+      if isinstance(arg, Tensor):
+        env = arg._env
+        break
+
+    with env:
+      return func(*args, **(kwargs or {}))
+
+  def detach(self):
+    return Tensor(jax.lax.stop_gradient(self.jax()), self._env)
+
+  def numpy(self) -> numpy.ndarray:
+    import numpy as np
+
+    return np.array(self._elem)
+
+  def jax(self) -> jax.Array:
+    return self._elem
+
+  def torch(self) -> torch.Tensor:
+    return j2t(self.jax())
+
+  @property
+  def dtype(self):
+    return j2t_dtype(self._elem.dtype)
+
+  def dim(self):
+    return self.ndim
+
+  @property
+  def device(self):
+    return torch.device('jax:0')
+
+  @property
+  def jax_device(self):
+    return self._elem.device
+
+  @property
+  def data(self):
+    logger.warn("In-place to .data modifications still results a copy on TPU")
+    return self
+
+  @data.setter
+  def data(self, other):
+    if isinstance(other, Tensor):
+      self._elem = other._elem
+
+  def apply_jax(self, jax_function, *args, **kwargs):
+    # Call a jax function on _elem
+    res = jax_function(self._elem, *args, **kwargs)
+    return self._env.j2t_iso(res)
+
+  def apply_jax_(self, jax_function, *args, **kwargs):
+    self._elem = jax_function(self._elem, *args, **kwargs)
+    return self
+
+  def tolist(self):
+    return self._elem.tolist()
+
+  def shard_(self, sharding):
+    self.apply_jax_(jax.lax.with_sharding_constraint, sharding)
+  
+
+def debug_accuracy(func, args, kwargs, current_output):
+  args_torch, kwargs_torch, out_torch = torch_pytree.tree_map_only(
+      torch.Tensor, lambda x: j2t(x._elem), (args, kwargs, current_output))
+
+  with mode_utils.no_dispatch(), torch._C.DisableTorchFunction():
+    if 'device' in kwargs_torch:
+      kwargs_torch['device'] = 'cpu'  # do the torch native for comparison
+    expected_out = func(*args_torch, **kwargs_torch)
+
+  flattened_current_out, _ = torch_pytree.tree_flatten(out_torch)
+  flattened_expected_out, _ = torch_pytree.tree_flatten(expected_out)
+
+  for ex, real in zip(flattened_expected_out, flattened_current_out):
+    if isinstance(ex, torch.Tensor) and ex.dtype != real.dtype:
+      ex = ex.to(real.dtype)
+    try:
+      if (isinstance(ex, torch.Tensor) and
+          not torch.allclose(ex, real, atol=1e-3, equal_nan=True)):
+        import pdb
+
+        pdb.set_trace()
+    except:
+      import pdb
+
+      pdb.set_trace()
+
+  return True
+
+def _make_debug_msg(is_dispatch, log_args, func, args, kwargs):
+  def _display(a):
+    if isinstance(a, torch.Tensor):
+      return f'Tensor of {type(a)}: {a.dtype}{a.shape}'
+    elif isinstance(a, jax.Array):
+      return f'Jax Array of {type(a)}: {a.dtype}{a.shape}'
+    else:
+      return str(a)
+
+  kwargs = kwargs or {}
+  title = 'DISPATCH' if is_dispatch else 'FUNCTION'
+  args_msg = 'args: ' + ','.join(_display(a) for a in args) if log_args else ''
+  kwargs_msg = 'kwargs: ' + ','.join(f'{key}: {_display(a)}' for key, a in kwargs.items()) if log_args else ''
+  return f'{title}: {_name_of_func(func)} {args_msg} ~ {kwargs_msg}'
+
+
+class XLAFunctionMode(torch.overrides.TorchFunctionMode):
+  """Context manager that dispatches torch function calls to JAX."""
+
+  def __init__(self, env):
+     self.env = env
+
+  def __torch_function__(self,
+                         func,
+                         types,
+                         args=(),
+                         kwargs=None) -> torch.Tensor:
+    message = f'FUNCTION: {_name_of_func(func)}'
+    if self.env.config.debug_print_each_op_operands:
+      message = message + 'f' 
+    message = _make_debug_msg(False, self.env.config.debug_print_each_op_operands,
+                              func, args, kwargs)
+    with log_nested(self.env, message):
+      try:
+        return self.env.dispatch(func, types, args, kwargs)
+      except OperatorNotFound:
+        pass
+      if _name_of_func(func) in ('rot90'): # skip rot90 with k%4==0 due to no change
+        if len(args) >= 2 and type(args[1]) == int:
+          if ((args[1])%4 == 0):
+            return args[0]
+      return func(*args, **(kwargs or {}))
+
+
+class XLADispatchMode(torch_dispatch.TorchDispatchMode):
+
+  def __init__(self, env):
+    self.env = env
+
+  def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+    message = _make_debug_msg(True, self.env.config.debug_print_each_op_operands,
+                              func, args, kwargs)
+    with log_nested(self.env, message):
+      if isinstance(func, torch._ops.OpOverloadPacket):
+        with self:
+          return func(*args, **kwargs)
+      if func.namespace not in ('aten', '_c10d_functional', 'torchvision'):
+        return func(*args, **kwargs)
+      return self.env.dispatch(func, types, args, kwargs)
+
+def _name_of_func(func):
+  if hasattr(func, 'name'):
+    return func.name()
+  return func.__name__
+
+
+# Constructors that don't take other tensor as input
+TENSOR_CONSTRUCTORS = {
+  torch.ones,
+  torch.zeros,
+  torch.empty,
+  torch.empty_strided,
+  torch.tensor,
+  torch.arange,
+  torch.eye,
+  torch.randn,
+  torch.rand,
+  torch.randint,
+  torch.full,
+  torch.as_tensor,
+}
+
+
+class Environment(contextlib.ContextDecorator):
+    """This class holds a set of configurations and "globals" needed
+
+    for executing torch program using jax.
+    Things included so far:
+
+    op registry
+    PRNGKey
+    Configs
+
+    Also helper functions to manipulate those.
+    """
+
+    _prng_key: jax.random.PRNGKey
+
+
+    def __init__(self, configuration=None):
+        self._function_mode = XLAFunctionMode(self)
+        self._dispatch_mode = XLADispatchMode(self)
+
+        # name is torch callable
+        self._ops = {}
+        self.load_ops()
+
+        self._mesh = None
+        self.config = configuration or config.Configuration()
+
+        self._manually_entered = False 
+        self.enabled = False
+        self._jax_devices = set(['jax', 'jax_cpu', 'xla'])
+
+    def get_as_jax_device(self, device: Any):
+      if device is None:
+        device = torch.get_default_device()
+
+      if isinstance(device, torch.device):
+        device = str(device)
+
+      if (not self.config.use_torch_native_for_cpu_tensor and 
+          device.startswith('cpu')):
+        return jax.devices('cpu')[0]
+
+      if self.config.treat_cuda_as_jax_device and device.startswith('cuda'):
+        return jax.local_devices()[0]
+
+      if device.startswith('jax'):
+        return jax.local_devices()[0]
+
+      return None # fallback to torch
+      
+
+
+    def load_ops(self):
+      from torchax.ops import jaten, jtorch, jc10d, jtorchvision_nms
+      self._ops.update(ops_registry.all_aten_ops)
+      self._ops.update(ops_registry.all_torch_functions)
+
+      decomps = torch._decomp.core_aten_decompositions()
+      from torchax.decompositions import EXTRA_DECOMP
+      decomps.update(EXTRA_DECOMP)
+      for k, v in decomps.items():
+        if k not in self._ops:
+          self._ops[k] = ops_registry.Operator(
+            k,
+            v,
+            is_jax_function=False,
+            is_user_defined=False,
+            needs_env=False
+          )
+
+    def _to_copy(self, the_tensor, new_dtype, new_device):
+      if isinstance(the_tensor, Tensor):
+        arr = the_tensor.jax()
+        if new_dtype is not None and new_dtype != arr.dtype:
+          arr = arr.astype(mappings.t2j_dtype(new_dtype))
+        if new_device is not None:
+          # convert xla tensor to other device
+          # only supported is CPU
+          if str(new_device).startswith('cpu'):
+            # converting to a non-jax device: let torch native handle it
+            torch_tensor = j2t(arr) if isinstance(the_tensor, Tensor) else arr
+            with mode_utils.no_dispatch(), torch._C.DisableTorchFunction():
+              return torch_tensor.to(new_device)
+      else:
+        if new_dtype is not None and new_dtype != the_tensor.dtype:
+          with mode_utils.no_dispatch(), torch._C.DisableTorchFunction():
+            the_tensor = the_tensor.to(new_dtype)
+        jax_device = self.get_as_jax_device(new_device)
+        if jax_device:
+          arr = t2j(the_tensor)
+          arr = jax.device_put(arr, jax_device)
+        else:
+          with mode_utils.no_dispatch(), torch._C.DisableTorchFunction():
+            return the_tensor.to(new_device)
+
+      return Tensor(arr, self)
+      
+
+    def get_and_rotate_prng_key(self, generator: Optional[torch.Generator]=None):
+      # Always use the default `randint` to get the next seed
+      with mode_utils.no_dispatch(), torch._C.DisableTorchFunction():
+        next_key = torch.randint(
+            0, 2**32, (), dtype=torch.uint32, generator=generator).numpy()
+
+      return jax.random.key(next_key)
+
+    def _handle_tensor_constructor(self, func, args, kwargs):
+      device = kwargs.get('device')
+      jax_device = self.get_as_jax_device(device)
+      # TODO(qihqi) figure out better ways for device propagation
+      if not self._manually_entered and jax_device is None:
+        # let torch handle it
+        with mode_utils.no_dispatch(), torch._C.DisableTorchFunction():
+          return func(*args, **kwargs)
+      with jax.default_device(jax_device):
+        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)
+        return res
+
+    def _torch_Tensor_to(self, args, kwargs):
+      the_tensor = args[0]
+      args = args[1:]
+      if len(args) >= 1 and isinstance(args[0], torch.Tensor):
+        dtype = args[0].dtype
+        device = args[0].device
+        return self._to_copy(the_tensor, dtype, device)
+      device = kwargs.get('device')
+      dtype = kwargs.get('dtype')
+      # args like pin_memory etc that we will ignore
+      args = list(filter(lambda x: not isinstance(x, bool), args))
+      if len(args) >= 2:
+        device, dtype, *_ = args
+      elif len(args) == 1 and isinstance(args[0], torch.dtype):
+        dtype = args[0]
+      elif len(args) == 1:
+        device = args[0]
+      return self._to_copy(the_tensor, dtype, device)
+
+
+    def dispatch(self, func, types, args, kwargs):
+
+      kwargs = kwargs or {}
+      if func in TENSOR_CONSTRUCTORS:
+        return self._handle_tensor_constructor(func, args, kwargs)
+      if func in (torch.Tensor.to, torch.ops.aten.lift_fresh.default ,torch.ops.aten._to_copy, torch.ops.aten._to_copy.default):
+        return self._torch_Tensor_to(args, kwargs)
+
+      # If the func doesn't act on Tensor, and is not a tensor constructor,
+      # We should skip and let torch handle it.
+      
+      tensor_args = [t for t in torch_pytree.tree_flatten(args)[0] if isinstance(t, torch.Tensor)]
+      if tensor_args and all(not isinstance(t, Tensor) for t in tensor_args):
+        return func(*args, **kwargs)
+
+      with jax.named_scope(_name_of_func(func)):
+        op = self._ops.get(func)
+
+        if op is None and isinstance(func, torch._ops.OpOverloadPacket):
+          op = self._ops.get(func.default)
+
+        if op is None and isinstance(func, torch._ops.OpOverload):
+          op = self._ops.get(func.overloadpacket)
+
+        if op is None:
+          raise OperatorNotFound(
+            f'Operator with name {_name_of_func(func)} has no lowering')
+
+        old_args, old_kwargs = args, kwargs
+        args, kwargs = torch_pytree.tree_map_only(
+            torch.distributed._functional_collectives.AsyncCollectiveTensor,
+            torch.distributed._functional_collectives.wait_tensor,
+            (args, kwargs))
+        try:
+          if op.is_jax_function:
+            args, kwargs = self.t2j_iso((args, kwargs))
+        except AssertionError:
+          if self.config.debug_mixed_tensor:
+            import pdb; pdb.set_trace()
+          else:
+            raise
+
+
+        if op.needs_env:
+          kwargs['env'] = self
+
+        with self:
+          res = op.func(*args, **kwargs)
+
+        if op.is_jax_function:
+          res = self.j2t_iso(res)
+
+        if self.config.debug_accuracy_for_each_op:
+          debug_accuracy(func, old_args, old_kwargs, res)
+        return res
+
+    def enable_torch_modes(self):
+      self._dispatch_mode.__enter__()
+      self._function_mode.__enter__()
+      self.enabled = True
+    
+    def disable_torch_modes(self, *exc):
+      if not exc:
+        exc = (None, None, None)
+      self._function_mode.__exit__(*exc)
+      self._dispatch_mode.__exit__(*exc)
+      self.enabled = False
+
+    def __enter__(self):
+      self.enable_torch_modes()
+      self._manually_entered = True
+      return self
+
+    def __exit__(self, *exc):
+      self._manually_entered = False
+      self.disable_torch_modes(*exc)
+
+    def _move_one_value(self, val):
+      if isinstance(val, torch.nn.Module):
+        with self:
+          return val.to('jax')
+      if isinstance(val, Tensor):
+        return val
+      if isinstance(val, torch.Tensor):
+        return Tensor(t2j(val), self)
+      return val
+
+    def to_xla(self, torchvalues):
+      # tensors are torch.Tensors (not XLATensor)
+      res = torch_pytree.tree_map(
+        self._move_one_value,
+          torchvalues)
+      return res
+
+    def t2j_iso(self, torchtensors):
+      def to_jax(x):
+        if isinstance(x, torch.distributed._functional_collectives.AsyncCollectiveTensor):
+          x = x.wait()
+        assert isinstance(x, Tensor), f'Expect a Tensor but got {type(x)}; usually this means there is a mixed math between XLATensor and torch.Tensor'
+        return x.jax()
+      return torch_pytree.tree_map_only(torch.Tensor, to_jax, torchtensors)
+
+    def j2t_iso(self, jaxarray):
+      return torch_pytree.tree_map_only(
+        jnp.ndarray, lambda x: Tensor(x, self), jaxarray)
+
+    def j2t_copy(self, args):
+      pass
+
+    def override_op_definition(self, op_to_override, op_impl):
+      self._ops[op_to_override] = ops_registry.Operator(
+        op_to_override,
+        op_impl,
+        is_jax_function=False,
+        is_user_defined=True,
+        needs_env=False
+      )

torchax/tf_integration.py

@@ -0,0 +1,119 @@
+# pylint: disable
+import os
+from typing import Any, Tuple
+
+from jax.experimental import jax2tf
+import tensorflow as tf
+import torch
+from torchax import export
+
+
+def exported_program_to_tf_function(ep, enable_xla=True):
+  weights, jax_program = export.exported_program_to_jax(ep)
+  wrapped = lambda *args: jax_program(weights, (args,))
+  avals = export.extract_avals(ep)
+  input_signature = [
+      tf.TensorSpec(shape=t.shape, dtype=t.dtype, name=f"args_{i}")
+      for i, t in enumerate(avals)
+  ]
+  tf_f = tf.function(
+      jax2tf.convert(
+          wrapped,
+          with_gradient=False,
+          enable_xla=enable_xla,
+      ),
+      autograph=False,
+      input_signature=input_signature,
+  )
+  return tf_f
+
+
+def exported_program_to_tf_module(ep: torch.export.ExportedProgram,
+                                  enable_xla=True) -> tf.Module:
+  tfm = tf.Module()
+  tfm.f = exported_program_to_tf_function(ep, enable_xla)
+  return tfm
+
+
+def save_exported_program_as_tf_saved_model(
+    ep: torch.export.ExportedProgram,
+    saved_model_dir: os.PathLike,
+    serving_key: str = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY,
+    function_alias: str = "",
+    enable_xla=True,
+):
+  """This function will export and save a pytorch ExportedProgram to tf.saved_model format.
+
+  The resulting tf.saved_model can be used inference using tf.serving model
+  server
+  or further convert to tflite flatbuffer for on-device serving.
+
+  Args:
+    torch_model: torch.nn.Module - model to export and save
+    args: Tuple[Any] - a set of args to trace the model with, i.e.
+      torch_model(*args) must run
+    saved_model_dir: os.PathLike - location to an empty directory to store the
+      saved_model
+    serving_key: str  - serving key tag, this is used by tf.serving to know
+      which function to run.
+    function_alias: str - passed through saved_model.save, used to tag a
+      function for inference converter or other tools.
+  """
+  tfm = exported_program_to_tf_module(ep, enable_xla=enable_xla)
+  signatures = {
+      serving_key: tfm.f.get_concrete_function(*tfm.f.input_signature)
+  }
+  save_options = tf.saved_model.SaveOptions(function_aliases={
+      function_alias: tfm.f,
+  })
+  tf.saved_model.save(
+      tfm,
+      saved_model_dir,
+      signatures=signatures,
+      options=save_options,
+  )
+
+
+def save_torch_module_as_tf_saved_model(
+    torch_model: torch.nn.Module,
+    args: Tuple[Any],
+    saved_model_dir: os.PathLike,
+    serving_key: str = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY,
+    function_alias: str = "",
+    enable_xla=True,
+):
+  """This function will export and save a pytorch nn.Module to tf.saved_model format.
+
+  The resulting tf.saved_model can be used inference using tf.serving model
+  server
+  or further convert to tflite flatbuffer for on-device serving.
+
+  Args:
+    torch_model: torch.nn.Module - model to export and save
+    args: Tuple[Any] - a set of args to trace the model with, i.e.
+      torch_model(*args) must run
+    saved_model_dir: os.PathLike - location to an empty directory to store the
+      saved_model
+    serving_key: str  - serving key tag, this is used by tf.serving to know
+      which function to run.
+    function_alias: str - passed through saved_model.save, used to tag a
+      function for inference converter or other tools.
+  """
+  ep = torch.export.export(torch_model, args)
+  save_exported_program_as_tf_saved_model(ep, saved_model_dir, serving_key,
+                                          function_alias, enable_xla)
+
+
+def exported_program_to_tflite_flatbuffer(ep: torch.export.ExportedProgram):
+  tfm = exported_program_to_tf_module(ep)
+  tf_concrete_func = tfm.f.get_concrete_function(*tfm.f.input_signature)
+  converter = tf.lite.TFLiteConverter.from_concrete_functions(
+      [tf_concrete_func], tfm)
+  tflite_model = converter.convert()
+  return tflite_model
+
+
+def torch_module_to_tflite_flatbuffer(torch_model: torch.nn.Module,
+                                      args: Tuple[Any]):
+  ep = torch.export.export(torch_model, args)
+  return exported_program_to_tflite_flatbuffer(ep)