torchax 0.0.10.dev20251118__py3-none-any.whl

torchax/CONTRIBUTING.md

@@ -0,0 +1,43 @@
+# Contributing to torchax
+
+We appreciate all contributions. If you are planning to contribute a bug fix for an open issue, please comment on the thread and we're happy to provide any guidance. You are very welcome to pick issues from good first issue and help wanted labels.
+
+
+# Developer setup
+
+## Mac setup:
+@qihqi
+
+I am able to develop directly on mac (m1) laptop for most of parts. Using steps
+in README.md works. The condensed version for easy copy & paste:
+
+```bash
+conda create --name <your_name> python=3.10
+conda activate <your_name>
+pip install --upgrade "jax[cpu]" torch
+pip install -r test_requirements.txt
+pip install -e .
+pip install pytest-xdist # recommended for running test faster
+pytest -n auto test
+```
+
+## Setup on GPU or TPU
+
+Same as Mac setup, except, if you run test using pytest, please also
+add `JAX_PLATFORMS=cpu`. The reason is because pytest usually runs
+test in multiple threads. CPU device can be accessed concurrently where
+TPU devices usually only allow one accesor per process; so it could deadlock.
+
+### VSCode
+
+I use vscode on my Mac. I loosely followed instruction in
+https://code.visualstudio.com/docs/python/python-tutorial
+to setup a proper python environment.
+
+The plugins I installed (a subset of the ones listed above) are:
+* VSCode's official Python plugin
+* Ruff formatter
+* Python Debugger
+
+I also changed Python interpreter to point at the one in my conda env.
+That is all the changes I have.

torchax/__init__.py

@@ -0,0 +1,149 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import contextlib
+import dataclasses
+import os
+from contextlib import contextmanager
+from typing import Any
+
+import jax
+import torch
+from torch.utils import _pytree as pytree
+
+import torchax.device_module
+from torchax import tensor
+
+from .checkpoint import load_checkpoint, save_checkpoint
+
+__version__ = "0.0.10.dev20251118"
+VERSION = __version__
+
+# the "fast path" uses some sparse tensor thingies that currently we
+# don't support
+torch.backends.mha.set_fastpath_enabled(False)
+
+
+__all__ = [
+  "default_env",
+  "extract_jax",
+  "enable_globally",
+  "save_checkpoint",
+  "load_checkpoint",
+]
+
+
+os.environ.setdefault("ENABLE_RUNTIME_UPTIME_TELEMETRY", "1")
+
+# torchax:oss-begin
+if getattr(jax.config, "jax_pjrt_client_create_options", None):
+  jax.config.update(
+    "jax_pjrt_client_create_options",
+    f"ml_framework_name:PyTorch/XLA2;ml_framework_version:{'v0.0.1'}",
+  )
+# torchax:oss-end
+
+env = None
+
+
+def default_env():
+  global env
+
+  if env is None:
+    env = tensor.Environment()
+  return env
+
+
+def extract_jax(mod: torch.nn.Module, env=None):
+  """Returns a pytree of jax.ndarray and a jax callable."""
+  if env is None:
+    env = default_env()
+  states = dict(mod.named_buffers())
+  states.update(mod.named_parameters())
+
+  states = env.t2j_copy(states)
+
+  # @jax.jit
+  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, args, kwargs, tie_weights=False)
+    return env.t2j_iso(res)
+
+  return states, jax_func
+
+
+def enable_globally():
+  env = default_env().enable_torch_modes()
+  return env
+
+
+def disable_globally():
+  global env
+  default_env().disable_torch_modes()
+
+
+@contextlib.contextmanager
+def disable_temporarily():
+  prev = default_env().enabled
+  if prev:
+    disable_globally()
+  yield ()
+  if prev:
+    enable_globally()
+
+
+torch.utils.rename_privateuse1_backend("jax")
+unsupported_dtype = [torch.quint8]
+
+
+torch._register_device_module("jax", torchax.device_module)
+
+
+def enable_accuracy_mode():
+  jax.config.update("jax_enable_x64", True)
+  jax.config.update("jax_default_matmul_precision", "highest")
+  default_env().config.internal_respect_torch_return_dtypes = True
+
+
+def enable_performance_mode():
+  jax.config.update("jax_enable_x64", False)
+  jax.config.update("jax_default_matmul_precision", "default")
+  default_env().config.internal_respect_torch_return_dtypes = False
+
+
+@dataclasses.dataclass
+class CompileOptions:
+  # only valid if compiling nn.Module
+  methods_to_compile: list[str] = dataclasses.field(default_factory=lambda: ["forward"])
+  jax_jit_kwargs: dict[str, Any] = dataclasses.field(default_factory=dict)
+  mode: str = "jax"  # or dynamo or export
+
+
+def compile(fn, options: CompileOptions | None = None):
+  options = options or CompileOptions()
+  if options.mode == "jax":
+    from torchax import interop
+
+    if isinstance(fn, torch.nn.Module):
+      module = interop.JittableModule(fn, extra_jit_args=options.jax_jit_kwargs)
+      for n in options.methods_to_compile:
+        module.make_jitted(n)
+      return module
+    else:
+      return interop.jax_jit(fn)
+  elif options.mode == "dynamo":
+    raise RuntimeError("dynamo mode is not supported yet")
+  elif options.mode == "export":
+    raise RuntimeError("export mode is not supported yet")

torchax/amp.py

@@ -0,0 +1,218 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import contextlib
+import enum
+
+import torch
+from torch.utils import _pytree as pytree
+
+
+# enum class CastPolicy : uint8_t {
+#   lower_precision_fp = 0, // Cast all inputs to lower_precision_fp before
+#                           // running the op. Currently, lower_precision_fp is
+#                           // fp16 for AutocastCUDA, and is defined by user
+#                           // (default bf16) for AutocastCPU or other device.
+#   fp32, // Cast all inputs to at::kFloat before running the op.
+#   fp32_set_opt_dtype, // Treats functions (like softmax) that
+#                       //  1. we'd like to run in fp32 and
+#                       //  2. have a std::optional<ScalarType> arg that controls
+#                       //  the output type.
+#                       // fp32_set_opt_dtype wrappers' policy is: if the output
+#                       // type is already set, don't touch it, otherwise, set
+#                       // it to at::kFloat.
+#   fp32_append_dtype, // Treats functions (like norm) that
+#                      //  1. we'd like to run in fp32 and
+#                      //  2. have some overloads that accept an output type and
+#                      //  other overloads that don't.
+#                      // fp32_append_dtype wrappers wrap the overloads that don't
+#                      // have an output dtype.
+#                      // The wrapper policy is:  append at::kFloat to the args,
+#                      // and redispatch to the type-aware overload.
+#   promote, // Run in the widest dtype among several args.
+# };
+class CastPolicy(enum.Enum):
+  LOWER_PRECISION_FP = 0
+  FP32 = 1
+  FP32_SET_OPT_DTYPE = 2
+  FP32_APPEND_DTYPE = 3
+  PROMOTE = 4
+
+
+def execute_policy(policy, args, kwargs, target_lower_fp):
+  def is_float(a):
+    return isinstance(a, torch.Tensor) and a.is_floating_point()
+
+  match policy:
+    case CastPolicy.LOWER_PRECISION_FP:
+      return pytree.tree_map_only(
+        is_float, lambda a: a.to(target_lower_fp), (args, kwargs)
+      )
+    case CastPolicy.FP32:
+      return pytree.tree_map_only(
+        is_float, lambda a: a.to(torch.float32), (args, kwargs)
+      )
+    case CastPolicy.PROMOTE:
+      dtypes = {a.dtype for a in args}
+      widest = max((dtype.itemsize, dtype) for dtype in dtypes)[1]
+      return pytree.tree_map_only(is_float, lambda a: a.to(widest), (args, kwargs))
+    case _:
+      raise AssertionError(f"Policy {policy} not implemented yet.")
+
+
+@contextlib.contextmanager
+def autocast(device, dtype=torch.bfloat16, env=None):
+  del device
+  if env is None:
+    import torchax
+
+    env = torchax.default_env()
+  with env.override_property(autocast_dtype=dtype):
+    yield
+
+
+# https://github.com/pytorch/pytorch/blob/05faba40287cf7d8734da96cb2e904f39710bf29/aten/src/ATen/autocast_mode.cpp#L327
+autocast_policy = {
+  torch.ops.aten.conv1d.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv1d.padding: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv2d.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv2d.padding: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv3d.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv3d.padding: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.bmm.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.mm.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.linalg_vecdot.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.baddbmm.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.addmm.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten._addmm_activation.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.addbmm.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.linear.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten._convolution.deprecated: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.matmul.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv_tbc.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.mkldnn_rnn_layer.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv_transpose1d.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv_transpose2d.input: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.conv_transpose3d.input: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.prelu.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten.scaled_dot_product_attention.default: CastPolicy.LOWER_PRECISION_FP,
+  torch.ops.aten._native_multi_head_attention.default: CastPolicy.LOWER_PRECISION_FP,
+  # fp32 cast policy
+  torch.ops.aten.avg_pool3d.default: CastPolicy.FP32,
+  torch.ops.aten.binary_cross_entropy.default: CastPolicy.FP32,
+  torch.ops.aten.grid_sampler.default: CastPolicy.FP32,
+  torch.ops.aten.polar.default: CastPolicy.FP32,
+  torch.ops.aten.prod.default: CastPolicy.FP32,
+  torch.ops.aten.prod.dim_int: CastPolicy.FP32,
+  torch.ops.aten.prod.dim_Dimname: CastPolicy.FP32,
+  torch.ops.aten.quantile.default: CastPolicy.FP32,
+  torch.ops.aten.quantile.scalar: CastPolicy.FP32,
+  torch.ops.aten.nanquantile.default: CastPolicy.FP32,
+  torch.ops.aten.nanquantile.scalar: CastPolicy.FP32,
+  torch.ops.aten.stft.default: CastPolicy.FP32,
+  torch.ops.aten.stft.center: CastPolicy.FP32,
+  torch.ops.aten.cdist.default: CastPolicy.FP32,
+  torch.ops.aten.grid_sampler_2d.default: CastPolicy.FP32,
+  torch.ops.aten._grid_sampler_2d_cpu_fallback.default: CastPolicy.FP32,
+  torch.ops.aten.grid_sampler_3d.default: CastPolicy.FP32,
+  torch.ops.aten.trace.default: CastPolicy.FP32,
+  torch.ops.aten.view_as_complex.default: CastPolicy.FP32,
+  torch.ops.aten.cholesky.default: CastPolicy.FP32,
+  torch.ops.aten.cholesky_inverse.default: CastPolicy.FP32,
+  torch.ops.aten.cholesky_solve.default: CastPolicy.FP32,
+  torch.ops.aten.inverse.default: CastPolicy.FP32,
+  torch.ops.aten.lu_solve.default: CastPolicy.FP32,
+  torch.ops.aten.orgqr.default: CastPolicy.FP32,
+  torch.ops.aten.ormqr.default: CastPolicy.FP32,
+  torch.ops.aten.pinverse.default: CastPolicy.FP32,
+  torch.ops.aten.max_pool3d.default: CastPolicy.FP32,
+  torch.ops.aten.max_unpool2d.default: CastPolicy.FP32,
+  torch.ops.aten.max_unpool3d.default: CastPolicy.FP32,
+  torch.ops.aten.adaptive_avg_pool3d.default: CastPolicy.FP32,
+  torch.ops.aten.reflection_pad1d.default: CastPolicy.FP32,
+  torch.ops.aten.reflection_pad2d.default: CastPolicy.FP32,
+  torch.ops.aten.replication_pad1d.default: CastPolicy.FP32,
+  torch.ops.aten.replication_pad2d.default: CastPolicy.FP32,
+  torch.ops.aten.replication_pad3d.default: CastPolicy.FP32,
+  torch.ops.aten.mse_loss.default: CastPolicy.FP32,
+  torch.ops.aten.cosine_embedding_loss.default: CastPolicy.FP32,
+  torch.ops.aten.nll_loss.default: CastPolicy.FP32,
+  torch.ops.aten.nll_loss2d.default: CastPolicy.FP32,
+  torch.ops.aten.hinge_embedding_loss.default: CastPolicy.FP32,
+  torch.ops.aten.poisson_nll_loss.default: CastPolicy.FP32,
+  torch.ops.aten.smooth_l1_loss.default: CastPolicy.FP32,
+  torch.ops.aten.cross_entropy_loss.default: CastPolicy.FP32,
+  torch.ops.aten.l1_loss.default: CastPolicy.FP32,
+  torch.ops.aten.huber_loss.default: CastPolicy.FP32,
+  torch.ops.aten.margin_ranking_loss.default: CastPolicy.FP32,
+  torch.ops.aten.soft_margin_loss.default: CastPolicy.FP32,
+  torch.ops.aten.triplet_margin_loss.default: CastPolicy.FP32,
+  torch.ops.aten.multi_margin_loss.default: CastPolicy.FP32,
+  torch.ops.aten.ctc_loss.IntList: CastPolicy.FP32,
+  torch.ops.aten.ctc_loss.Tensor: CastPolicy.FP32,
+  torch.ops.aten.kl_div.default: CastPolicy.FP32,
+  torch.ops.aten.multilabel_margin_loss.default: CastPolicy.FP32,
+  torch.ops.aten.binary_cross_entropy_with_logits.default: CastPolicy.FP32,
+  torch.ops.aten.fft_fft.default: CastPolicy.FP32,
+  torch.ops.aten.fft_ifft.default: CastPolicy.FP32,
+  torch.ops.aten.fft_fft2.default: CastPolicy.FP32,
+  torch.ops.aten.fft_ifft2.default: CastPolicy.FP32,
+  torch.ops.aten.fft_fftn.default: CastPolicy.FP32,
+  torch.ops.aten.fft_ifftn.default: CastPolicy.FP32,
+  torch.ops.aten.fft_rfft.default: CastPolicy.FP32,
+  torch.ops.aten.fft_irfft.default: CastPolicy.FP32,
+  torch.ops.aten.fft_rfft2.default: CastPolicy.FP32,
+  torch.ops.aten.fft_irfft2.default: CastPolicy.FP32,
+  torch.ops.aten.fft_rfftn.default: CastPolicy.FP32,
+  torch.ops.aten.fft_irfftn.default: CastPolicy.FP32,
+  torch.ops.aten.fft_hfft.default: CastPolicy.FP32,
+  torch.ops.aten.fft_ihfft.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_cond.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_cond.p_str: CastPolicy.FP32,
+  torch.ops.aten.linalg_matrix_rank.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_matrix_rank.tol_tensor: CastPolicy.FP32,
+  torch.ops.aten.linalg_matrix_rank.atol_rtol_tensor: CastPolicy.FP32,
+  torch.ops.aten.linalg_matrix_rank.atol_rtol_float: CastPolicy.FP32,
+  torch.ops.aten.linalg_solve.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_cholesky.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_svdvals.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_eigvals.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_eigvalsh.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_inv.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_householder_product.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_tensorinv.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_tensorsolve.default: CastPolicy.FP32,
+  torch.ops.aten.fake_quantize_per_tensor_affine.default: CastPolicy.FP32,
+  torch.ops.aten.geqrf.default: CastPolicy.FP32,
+  torch.ops.aten._lu_with_info.default: CastPolicy.FP32,
+  torch.ops.aten.qr.default: CastPolicy.FP32,
+  torch.ops.aten.svd.default: CastPolicy.FP32,
+  torch.ops.aten.triangular_solve.default: CastPolicy.FP32,
+  torch.ops.aten.fractional_max_pool2d.default: CastPolicy.FP32,
+  torch.ops.aten.fractional_max_pool3d.default: CastPolicy.FP32,
+  torch.ops.aten.adaptive_max_pool3d.default: CastPolicy.FP32,
+  torch.ops.aten.multilabel_margin_loss_forward.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_qr.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_cholesky_ex.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_svd.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_eig.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_eigh.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_lstsq.default: CastPolicy.FP32,
+  torch.ops.aten.linalg_inv_ex.default: CastPolicy.FP32,
+  # promote
+  torch.ops.aten.stack.default: CastPolicy.PROMOTE,
+  torch.ops.aten.cat.default: CastPolicy.PROMOTE,
+  torch.ops.aten.index_copy.default: CastPolicy.PROMOTE,
+  torch.ops.aten.index_copy.dimname: CastPolicy.PROMOTE,
+}

torchax/checkpoint.py

@@ -0,0 +1,85 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import torch
+from flax.training import checkpoints
+
+from . import tensor
+
+
+def _to_jax(pytree):
+  def to_jax_array(x):
+    if isinstance(x, tensor.Tensor):
+      return x.jax()
+    elif isinstance(x, torch.Tensor):
+      return jnp.asarray(x.cpu().numpy())
+    return x
+
+  return jax.tree_util.tree_map(to_jax_array, pytree)
+
+
+def _to_torch(pytree):
+  return jax.tree_util.tree_map(
+    lambda x: torch.from_numpy(np.asarray(x))
+    if isinstance(x, (jnp.ndarray, jax.Array))
+    else x,
+    pytree,
+  )
+
+
+def save_checkpoint(state: dict[str, Any], path: str, step: int):
+  """Saves a checkpoint to a file in JAX style.
+
+  Args:
+    state: A dictionary containing the state to save. torch.Tensors will be
+      converted to jax.Array.
+    path: The path to save the checkpoint to. This is a directory.
+    step: The training step.
+  """
+  state = _to_jax(state)
+  checkpoints.save_checkpoint(path, state, step=step, overwrite=True)
+
+
+def load_checkpoint(path: str) -> dict[str, Any]:
+  """Loads a checkpoint and returns it in JAX format.
+
+  This function can load both PyTorch-style (single file) and JAX-style
+  (directory) checkpoints.
+
+  If the checkpoint is in PyTorch format, it will be converted to JAX format.
+
+  Args:
+    path: The path to the checkpoint.
+
+  Returns:
+    The loaded state in JAX format (pytree with jax.Array leaves).
+  """
+  if os.path.isdir(path):
+    # JAX-style checkpoint
+    state = checkpoints.restore_checkpoint(path, target=None)
+    if state is None:
+      raise FileNotFoundError(f"No checkpoint found at {path}")
+    return state
+  elif os.path.isfile(path):
+    # PyTorch-style checkpoint
+    state = torch.load(path, weights_only=False)
+    return _to_jax(state)
+  else:
+    raise FileNotFoundError(f"No such file or directory: {path}")

torchax/config.py

@@ -0,0 +1,44 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+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