torchax 0.0.6__py3-none-any.whl → 0.0.10.dev20251116__py3-none-any.whl

torchax/CONTRIBUTING.md

@@ -1,9 +1,7 @@
-# Contributing to TorchXLA2
+# 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.
 
-If you plan to contribute new features, utility functions or extensions to the core, please first open an issue and discuss the feature with us. Sending a PR without discussion might end up resulting in a rejected PR, because we might be taking the core in a different direction than you might be aware of.
-
 
 # Developer setup
 
@@ -19,9 +17,17 @@ conda activate <your_name>
 pip install --upgrade "jax[cpu]" torch
 pip install -r test_requirements.txt
 pip install -e .
-pytest test
+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
@@ -35,4 +41,3 @@ The plugins I installed (a subset of the ones listed above) are:
 
 I also changed Python interpreter to point at the one in my conda env.
 That is all the changes I have.
-

torchax/__init__.py

@@ -1,3 +1,17 @@
+# 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
 from typing import List, Dict, Any, Optional
 import dataclasses
@@ -8,119 +22,132 @@ from torch.utils import _pytree as pytree
 from torchax import tensor
 from contextlib import contextmanager
 
-__version__ = "0.0.6"
+__version__ = "0.0.10.dev20251116"
 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',
+    "default_env",
+    "extract_jax",
+    "enable_globally",
+    "save_checkpoint",
+    "load_checkpoint",
 ]
 
-from jax._src import xla_bridge
+from .checkpoint import save_checkpoint, load_checkpoint
 
-os.environ.setdefault('ENABLE_RUNTIME_UPTIME_TELEMETRY', '1')
+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"}')
+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
+    global env
 
-  if env is None:
-    env = tensor.Environment()
-  return 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())
+    """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)
+    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)
+    # @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
+    return states, jax_func
 
 
 def enable_globally():
-  env = default_env().enable_torch_modes()
-  return env
+    env = default_env().enable_torch_modes()
+    return env
 
 
 def disable_globally():
-  global env
-  default_env().disable_torch_modes()
+    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()
+    prev = default_env().enabled
+    if prev:
+        disable_globally()
+    yield ()
+    if prev:
+        enable_globally()
 
 
-torch.utils.rename_privateuse1_backend('jax')
+torch.utils.rename_privateuse1_backend("jax")
 unsupported_dtype = [torch.quint8]
 
 import jax
 import torchax.device_module
 
-torch._register_device_module('jax', torchax.device_module)
+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
+    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
+    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
+    # 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: Optional[CompileOptions] = 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')
+    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

@@ -1,3 +1,17 @@
+# 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

torchax/checkpoint.py

@@ -0,0 +1,79 @@
+# 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 torch
+import os
+from typing import Any, Dict
+from flax.training import checkpoints
+import jax
+import jax.numpy as jnp
+import numpy as np
+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

@@ -1,3 +1,17 @@
+# 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
 
 

torchax/decompositions.py

@@ -1,3 +1,17 @@
+# 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.
+
 """This file contains some decompositons that are not available in torch stable.
 
 Most likely from Content of

torchax/device_module.py

@@ -1,3 +1,17 @@
+# 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 torch
 
 

torchax/export.py

@@ -1,3 +1,17 @@
+# 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.
+
 # pylint: disable
 """Utilities for exporting a torch program to jax/stablehlo."""
 import copy

torchax/flax.py

@@ -1,3 +1,17 @@
+# 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.
+
 """Flax interop."""
 
 import torch

torchax/interop.py

@@ -1,3 +1,17 @@
+# 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 collections
 import copy
 import functools
@@ -182,7 +196,6 @@ def _torch_view(t: JaxValue) -> TorchValue:
   # t is an object from jax land
   # view it as-if it's a torch land object
   if isinstance(t, jax.Array):
-    # TODO
     return tensor.Tensor(t, torchax.default_env())
   if isinstance(t, jnp.dtype):
     return mappings.j2t_dtype(t)
@@ -237,6 +250,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 +333,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/mesh_util.py

@@ -1,3 +1,17 @@
+# 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 jax
 from jax.sharding import PartitionSpec, NamedSharding
 import torch

torchax/ops/__init__.py

@@ -1,3 +1,17 @@
+# 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.
+
 def all_aten_jax_ops():
   # to load the ops
   import torchax.ops.jaten  # type: ignore