torchax 0.0.10.dev20251116__py3-none-any.whl → 0.0.11.dev202612__py3-none-any.whl

torchax/checkpoint.py

@@ -12,32 +12,39 @@
 # 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
+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()
+      return x.jax()
     elif isinstance(x, torch.Tensor):
-        return jnp.asarray(x.cpu().numpy())
+      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)
+    if isinstance(x, (jnp.ndarray, jax.Array))
+    else x,
+    pytree,
+  )
 
 
-def save_checkpoint(state: Dict[str, Any], path: str, step: int):
+def save_checkpoint(state: dict[str, Any], path: str, step: int):
   """Saves a checkpoint to a file in JAX style.
 
   Args:
@@ -50,7 +57,7 @@ def save_checkpoint(state: Dict[str, Any], path: str, step: int):
   checkpoints.save_checkpoint(path, state, step=step, overwrite=True)
 
 
-def load_checkpoint(path: str) -> Dict[str, Any]:
+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
@@ -76,4 +83,3 @@ def load_checkpoint(path: str) -> Dict[str, Any]:
     return _to_jax(state)
   else:
     raise FileNotFoundError(f"No such file or directory: {path}")
-    

torchax/config.py

@@ -35,10 +35,6 @@ class Configuration:
   # 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/decompositions.py

@@ -22,21 +22,21 @@ Can also contain decompositions of a torch op in terms of other torch ops.
 """
 
 import functools
-from typing import Any, Callable, List, Tuple
+from collections.abc import Callable
+from typing import Any
 
 import torch
-from torch import Tensor
 import torch._decomp as decomp
-from torch._decomp import decompositions_for_rng
-from torch._decomp import register_decomposition
 import torch._prims_common as utils
+from torch import Tensor
+from torch._decomp import decompositions_for_rng, register_decomposition
 from torch._prims_common.wrappers import out_wrapper
 
 DispatchKey = torch._C.DispatchKey  # type: ignore[attr-defined]
 
 # None of these functions are publicly accessible; get at them
 # from torch._decomps
-__all__: List[str] = []
+__all__: list[str] = []
 
 aten = torch._ops.ops.aten
 
@@ -44,21 +44,21 @@ aten = torch._ops.ops.aten
 def _try_register(op, impl):
   try:
     register_decomposition(op)(impl)
-  except:
+
+  except Exception:
     pass
 
 
 @out_wrapper()
-def _reflection_pad(a: Tensor, padding: Tuple[int, ...]) -> Tensor:
-
+def _reflection_pad(a: Tensor, padding: tuple[int, ...]) -> Tensor:
   def idx(left, middle, right):
     dim_idx = torch.arange(-left, middle + right, device=a.device)
     return middle - 1 - (middle - 1 - dim_idx.abs()).abs()
 
   return _reflection_or_replication_pad(
-      a,
-      padding,
-      idx,
+    a,
+    padding,
+    idx,
   )
 
 
@@ -68,32 +68,32 @@ _try_register(aten.reflection_pad3d, _reflection_pad)
 
 
 @out_wrapper()
-def _replication_pad(a: Tensor, padding: Tuple[int, ...]) -> Tensor:
-
+def _replication_pad(a: Tensor, padding: tuple[int, ...]) -> Tensor:
   def idx(left, middle, right):
     dim_idx = torch.arange(-left, middle + right, device=a.device)
     return torch.clamp(dim_idx, 0, middle - 1)
 
   return _reflection_or_replication_pad(
-      a,
-      padding,
-      idx,
+    a,
+    padding,
+    idx,
   )
 
 
-decomp.global_decomposition_table["post_autograd"][
-    aten.replication_pad2d.default] = _replication_pad
+decomp.global_decomposition_table["post_autograd"][aten.replication_pad2d.default] = (
+  _replication_pad
+)
 
 
 def _reflection_or_replication_pad(
-    a: Tensor,
-    padding: Tuple[int, ...],
-    idx_fn: Callable[[int, int, int], Tensor],
+  a: Tensor,
+  padding: tuple[int, ...],
+  idx_fn: Callable[[int, int, int], Tensor],
 ) -> Tensor:
   dim = len(padding) // 2
   torch._check(
-      a.dim() in (dim + 1, dim + 2),
-      lambda: f"reflection_pad{dim}d requires {dim + 1}D or {dim + 2}D input",
+    a.dim() in (dim + 1, dim + 2),
+    lambda: f"reflection_pad{dim}d requires {dim + 1}D or {dim + 2}D input",
   )
   inp_shape = a.shape[-dim:]
   nc_dim = a.dim() - dim
@@ -103,7 +103,7 @@ def _reflection_or_replication_pad(
 
   result = a
   for i in range(dim):
-    idx: List[Any] = [None] * result.dim()
+    idx: list[Any] = [None] * result.dim()
     idx[i + nc_dim] = idx_fn(padding_left[i], inp_shape[i], padding_right[i])
     result = aten._unsafe_index(result, idx)
 
@@ -158,11 +158,11 @@ def _sum_tensors(ts) -> Tensor:
 
 @register_decomposition(aten.grid_sampler_3d)
 def _grid_sampler_3d(
-    a: torch.Tensor,
-    grid: torch.Tensor,
-    interpolation_mode: int = 0,
-    padding_mode: int = 0,
-    align_corners: bool = False,
+  a: torch.Tensor,
+  grid: torch.Tensor,
+  interpolation_mode: int = 0,
+  padding_mode: int = 0,
+  align_corners: bool = False,
 ) -> Tensor:
   """References: https://github.com/pytorch/pytorch/blob/06a7dc21c1005750598c37f3adbc031183c74de6/torch/_decomp/decompositions.py#L4075
 
@@ -170,11 +170,12 @@ def _grid_sampler_3d(
   """
   _expand_grid = False
   torch._check(
-      interpolation_mode in (0, 1),
-      lambda: f"Invalid interpolation mode {interpolation_mode}",
+    interpolation_mode in (0, 1),
+    lambda: f"Invalid interpolation mode {interpolation_mode}",
   )
   torch._check(
-      padding_mode in (0, 1, 2), lambda: f"Invalid padding mode {padding_mode}")
+    padding_mode in (0, 1, 2), lambda: f"Invalid padding mode {padding_mode}"
+  )
 
   # a is 5D: [B, C, D, H, W]
 
@@ -189,8 +190,7 @@ def _grid_sampler_3d(
   # Reflects coordinates until they fall between low and high (inclusive).
   # The bounds are passed as twice their value so that half-integer values
   # can be represented as ints.
-  def reflect_coordinates(coords: Tensor, twice_low: int,
-                          twice_high: int) -> Tensor:
+  def reflect_coordinates(coords: Tensor, twice_low: int, twice_high: int) -> Tensor:
     if twice_low == twice_high:
       return torch.zeros_like(coords)
     coords_min = twice_low / 2
@@ -198,8 +198,9 @@ def _grid_sampler_3d(
     coords2 = (coords - coords_min).abs()
     extra = torch.fmod(coords2, coords_span)
     flips = (coords2 / coords_span).floor().to(dtype=torch.int8)
-    return torch.where(flips & 1 == 0, extra + coords_min,
-                       coords_span + coords_min - extra)
+    return torch.where(
+      flips & 1 == 0, extra + coords_min, coords_span + coords_min - extra
+    )
 
   def compute_coordinates(coords: Tensor, size: int) -> Tensor:
     if padding_mode == 0:  # Zero
@@ -219,7 +220,7 @@ def _grid_sampler_3d(
 
   N, C, iD, iH, iW = a.shape
   _, oD, oH, oW, three = grid.shape
-  assert three == 3, "Last dim of grid must be 3. got {}".format(three)
+  assert three == 3, f"Last dim of grid must be 3. got {three}"
 
   def in_bounds_cond(xs: Tensor, ys: Tensor, zs) -> Tensor:
     xcheck = torch.logical_and(0 <= xs, xs < iW)
@@ -238,15 +239,16 @@ def _grid_sampler_3d(
     # broadcasting with N_idx, C_idx for the purposes of advanced indexing
     c = C if _expand_grid else 1
     return tuple(
-        torch.where(cond, t, 0).view(N, c, oD, oH, oW) for t in (
-            xs.to(dtype=torch.int64),
-            ys.to(dtype=torch.int64),
-            zs.to(dtype=torch.int64),
-            ws,
-        ))
-
-  def get_summand(ix: torch.Tensor, iy: torch.Tensor, iz: torch.Tensor,
-                  w) -> Tensor:
+      torch.where(cond, t, 0).view(N, c, oD, oH, oW)
+      for t in (
+        xs.to(dtype=torch.int64),
+        ys.to(dtype=torch.int64),
+        zs.to(dtype=torch.int64),
+        ws,
+      )
+    )
+
+  def get_summand(ix: torch.Tensor, iy: torch.Tensor, iz: torch.Tensor, w) -> Tensor:
     # Perform clipping, index into input tensor and multiply by weight
     idx_x, idx_y, idx_z, w_ = clip(ix, iy, iz, w)
     return a[N_idx, C_idx, idx_z, idx_y, idx_x] * w_
@@ -279,16 +281,18 @@ def _grid_sampler_3d(
     w_seb = (ix - ix_nwf) * (iy - iy_nwf) * (id_ - id_nwf)
 
     return _sum_tensors(
-        get_summand(ix, iy, id_, w) for (ix, iy, id_, w) in (
-            (ix_nwf, iy_nwf, id_nwf, w_nwf),
-            (ix_nef, iy_nef, id_nef, w_nef),
-            (ix_swf, iy_swf, id_swf, w_swf),
-            (ix_sef, iy_sef, id_sef, w_sef),
-            (ix_nwb, iy_nwb, id_nwb, w_nwb),
-            (ix_neb, iy_neb, id_neb, w_neb),
-            (ix_swb, iy_swb, id_swb, w_swb),
-            (ix_seb, iy_seb, id_seb, w_seb),
-        ))
+      get_summand(ix, iy, id_, w)
+      for (ix, iy, id_, w) in (
+        (ix_nwf, iy_nwf, id_nwf, w_nwf),
+        (ix_nef, iy_nef, id_nef, w_nef),
+        (ix_swf, iy_swf, id_swf, w_swf),
+        (ix_sef, iy_sef, id_sef, w_sef),
+        (ix_nwb, iy_nwb, id_nwb, w_nwb),
+        (ix_neb, iy_neb, id_neb, w_neb),
+        (ix_swb, iy_swb, id_swb, w_swb),
+        (ix_seb, iy_seb, id_seb, w_seb),
+      )
+    )
   else:  # interpolation_mode == 1:  # Nearest
     ix = compute_source_index(x, iW)
     iy = compute_source_index(y, iH)
@@ -301,7 +305,8 @@ def _grid_sampler_3d(
     return get_summand(ix_nearest, iy_nearest, iz_nearest, 1)
 
 
-DECOMPOSITIONS = decomp.get_decompositions([
+DECOMPOSITIONS = decomp.get_decompositions(
+  [
     torch.ops.aten.upsample_bicubic2d,
     torch.ops.aten.upsample_nearest1d,
     torch.ops.aten.upsample_nearest2d,
@@ -782,9 +787,10 @@ DECOMPOSITIONS = decomp.get_decompositions([
     torch.ops.aten.__irshift__.Tensor,
     torch.ops.aten.__irshift__.Scalar,
     torch.ops.aten.__ior__.Tensor,
-])
+  ]
+)
 
 MUTABLE_DECOMPOSITION = [
-    torch.ops.aten.bernoulli_.Tensor,
-    torch.ops.aten.bernoulli_.float,
+  torch.ops.aten.bernoulli_.Tensor,
+  torch.ops.aten.bernoulli_.float,
 ]

torchax/export.py

@@ -14,17 +14,21 @@
 
 # pylint: disable
 """Utilities for exporting a torch program to jax/stablehlo."""
+
 import copy
-from typing import Any, Dict, Tuple
+from typing import Any
+
+import jax
+import jax.export
+import sympy
 import torch
+import torch._refs
+from torch._decomp import get_decompositions
 from torch.utils import _pytree as pytree
+
 import torchax
-from torchax import tensor
-from torchax.ops import ops_registry, mappings
 from torchax import decompositions
-import jax
-import jax.export
-import sympy
+from torchax.ops import mappings, ops_registry
 
 DEBUG = False
 
@@ -34,16 +38,13 @@ class JaxInterpreter(torch.fx.Interpreter):
 
   def __init__(self, graph_module):
     super().__init__(graph_module)
-    import torchax.ops.jaten
-    import torchax.ops.jtorch
 
-  def call_function(self, target, args: Tuple, kwargs: Dict) -> Any:
-    if not isinstance(target,
-                      (torch._ops.OpOverloadPacket, torch._ops.OpOverload)):
+  def call_function(self, target, args: tuple, kwargs: dict) -> Any:
+    if not isinstance(target, (torch._ops.OpOverloadPacket, torch._ops.OpOverload)):
       return super().call_function(target, args, kwargs)
 
     if DEBUG:
-      print('Running ', target.name(), '--------')
+      print("Running ", target.name(), "--------")
 
     op = ops_registry.all_aten_ops.get(target)
     if op is None:
@@ -54,31 +55,30 @@ class JaxInterpreter(torch.fx.Interpreter):
       op = ops_registry.all_aten_ops.get(target.overloadpacket)
     if op is None:
       print(target.name(), target.tags)
-      raise RuntimeError('No lowering found for', target.name())
+      raise RuntimeError("No lowering found for", target.name())
     return op.func(*args, **kwargs)
 
   def run_node(self, n) -> Any:
     res = super().run_node(n)
     if DEBUG:
-      if n.op == 'call_function':
-        if hasattr(res, 'shape'):
-          print('Meta:', n.meta.get('val').shape, 'REAL: ', res.shape)
+      if n.op == "call_function":
+        if hasattr(res, "shape"):
+          print("Meta:", n.meta.get("val").shape, "REAL: ", res.shape)
     return res
 
 
-from torch._decomp import get_decompositions
-import torch._refs
-
 _extra_decomp = get_decompositions([torch.ops.aten.unfold])
 
 
 def _extract_states_from_exported_program(exported_model):
   # NOTE call convention: (parameters, buffers, user_inputs)
-  param_and_buffer_keys = exported_model.graph_signature.parameters + exported_model.graph_signature.buffers
+  param_and_buffer_keys = (
+    exported_model.graph_signature.parameters + exported_model.graph_signature.buffers
+  )
   state_dict = copy.copy(exported_model.state_dict)
-  if (constants := getattr(exported_model, 'constants', None)) is not None:
+  if (constants := getattr(exported_model, "constants", None)) is not None:
     state_dict.update(constants)
-  param_buffer_values = list(state_dict[key] for key in param_and_buffer_keys)
+  param_buffer_values = [state_dict[key] for key in param_and_buffer_keys]
 
   if hasattr(exported_model.graph_signature, "lifted_tensor_constants"):
     for name in exported_model.graph_signature.lifted_tensor_constants:
@@ -94,19 +94,19 @@ def exported_program_to_jax(exported_program, export_raw: bool = False):
 
   func(state, input) would be how you call it.
   """
-  if torch.__version__ >= '2.2':
+  if torch.__version__ >= "2.2":
     # torch version 2.1 didn't expose this yet
     exported_program = exported_program.run_decompositions()
     exported_program = exported_program.run_decompositions(
-        decompositions.DECOMPOSITIONS)
+      decompositions.DECOMPOSITIONS
+    )
   if DEBUG:
     print(exported_program.graph_module.code)
 
   names, states = _extract_states_from_exported_program(exported_program)
 
   def _extract_args(args, kwargs):
-    flat_args, received_spec = pytree.tree_flatten(
-        (args, kwargs))  # type: ignore[possibly-undefined]
+    flat_args, received_spec = pytree.tree_flatten((args, kwargs))  # type: ignore[possibly-undefined]
     return flat_args
 
   num_mutations = len(exported_program.graph_signature.buffers_to_mutate)
@@ -114,9 +114,9 @@ def exported_program_to_jax(exported_program, export_raw: bool = False):
   def func(states, inputs):
     args = _extract_args(inputs, {})
     res = JaxInterpreter(exported_program.graph_module).run(
-        *states,
-        *args,
-        enable_io_processing=False,
+      *states,
+      *args,
+      enable_io_processing=False,
     )
     res = res[num_mutations:]
     return res
@@ -134,21 +134,19 @@ def extract_avals(exported):
   """
 
   def _to_aval(arg_meta, symbolic_shapes):
-    """Convet from torch type to jax abstract value for export tracing
-    """
+    """Convet from torch type to jax abstract value for export tracing"""
 
     def _get_dim(d):
       if isinstance(d, torch.SymInt):
         return symbolic_shapes[str(d)]
       return d
 
-    val = arg_meta['val']
-    is_scalar = isinstance(val, float) or isinstance(val, int) or isinstance(
-        val, bool)
+    val = arg_meta["val"]
+    is_scalar = isinstance(val, float) or isinstance(val, int) or isinstance(val, bool)
     if is_scalar:
-      return jax.ShapeDtypeStruct([], type(arg_meta['val']))
+      return jax.ShapeDtypeStruct([], type(arg_meta["val"]))
 
-    tensor_meta = arg_meta['tensor_meta']
+    tensor_meta = arg_meta["tensor_meta"]
     shape = [_get_dim(d) for d in tensor_meta.shape]
     return jax.ShapeDtypeStruct(shape, mappings.t2j_dtype(tensor_meta.dtype))
 
@@ -156,8 +154,9 @@ def extract_avals(exported):
     """Return placeholders with input metadata"""
     placeholders = [p for p in exported.graph.nodes if p.op == "placeholder"]
     input_placeholders = [
-        p for p, s in zip(placeholders, exported.graph_signature.input_specs)
-        if s.kind == torch.export.graph_signature.InputKind.USER_INPUT
+      p
+      for p, s in zip(placeholders, exported.graph_signature.input_specs, strict=False)
+      if s.kind == torch.export.graph_signature.InputKind.USER_INPUT
     ]
     return input_placeholders
 
@@ -179,17 +178,22 @@ def extract_avals(exported):
         torch.export.Dim("a", min=5, max=10)
          ==> ("a >= 5", "a <= 10",)
       """
-      if not isinstance(torch_constraint, torch.utils._sympy.value_ranges.
-                        ValueRanges) or torch_constraint.is_bool:
-        raise TypeError(
-            f"No symbolic constraint handler for: {torch_constraint}")
+      if (
+        not isinstance(torch_constraint, torch.utils._sympy.value_ranges.ValueRanges)
+        or torch_constraint.is_bool
+      ):
+        raise TypeError(f"No symbolic constraint handler for: {torch_constraint}")
 
       constraints = []
       symbol = sympy.Symbol(symbol_name)
       if torch_constraint.lower != 2:
         constraints.append(symbol >= torch_constraint.lower)
       from sympy.core.singleton import S
-      if not torch_constraint.upper.is_infinite and torch_constraint.upper is not S.IntInfinity:
+
+      if (
+        not torch_constraint.upper.is_infinite
+        and torch_constraint.upper is not S.IntInfinity
+      ):
         constraints.append(symbol <= torch_constraint.upper)
 
       return tuple(sympy.pretty(c, use_unicode=False) for c in constraints)
@@ -208,8 +212,7 @@ def extract_avals(exported):
       symbol_name = str(sym)
       constraints = _build_symbolic_constraints(symbol_name, constraint)
       if sym.is_symbol:
-        symbolic_shape = jax.export.symbolic_shape(
-            symbol_name, constraints=constraints)
+        symbolic_shape = jax.export.symbolic_shape(symbol_name, constraints=constraints)
       else:
         assert len(sym.free_symbols) > 0
         scope = free_symbols[str(list(sym.free_symbols)[0])].scope
@@ -222,12 +225,8 @@ def extract_avals(exported):
     # integer compuations on symbol variables, so each symbol variable is OK to
     # have its own scope.
     symbolic_shapes = {}
-    symbol_variables = [
-        (s, v) for s, v in range_constraints.items() if s.is_symbol
-    ]
-    symbol_exprs = [
-        (s, v) for s, v in range_constraints.items() if not s.is_symbol
-    ]
+    symbol_variables = [(s, v) for s, v in range_constraints.items() if s.is_symbol]
+    symbol_exprs = [(s, v) for s, v in range_constraints.items() if not s.is_symbol]
     for sym, constraint in symbol_variables + symbol_exprs:
       symbolic_shape = _build_symbolic_shape(sym, constraint, symbolic_shapes)
       symbolic_shapes[str(sym)] = symbolic_shape
@@ -237,10 +236,10 @@ def extract_avals(exported):
   args = _get_inputs(exported)
 
   if DEBUG:
-    print('Inputs to aval:', args, '--------')
-    print('Symbolic shapes:', symbolic_shapes)
+    print("Inputs to aval:", args, "--------")
+    print("Symbolic shapes:", symbolic_shapes)
     for arg in args:
-      print('Meta2Aval', arg.meta, '--> ', _to_aval(arg.meta, symbolic_shapes))
+      print("Meta2Aval", arg.meta, "--> ", _to_aval(arg.meta, symbolic_shapes))
 
   return [_to_aval(arg.meta, symbolic_shapes) for arg in args]
 

torchax/flax.py

@@ -15,18 +15,19 @@
 """Flax interop."""
 
 import torch
+
 import torchax as tx
 import torchax.interop
 
 
 class FlaxNNModule(torch.nn.Module):
-
   def __init__(self, env, flax_module, sample_args, sample_kwargs=None):
     super().__init__()
     prng = env.prng_key
     sample_kwargs = sample_kwargs or {}
-    parameter_dict = tx.interop.call_jax(flax_module.init, prng, *sample_args,
-                                         **sample_kwargs)
+    parameter_dict = tx.interop.call_jax(
+      flax_module.init, prng, *sample_args, **sample_kwargs
+    )
 
     self._params = self._encode_nested_dict(parameter_dict)
 
@@ -49,5 +50,6 @@ class FlaxNNModule(torch.nn.Module):
 
   def forward(self, *args, **kwargs):
     nested_dict_params = self._decode_nested_dict(self._params)
-    return tx.interop.call_jax(self._flax_module.apply, nested_dict_params,
-                               *args, **kwargs)
+    return tx.interop.call_jax(
+      self._flax_module.apply, nested_dict_params, *args, **kwargs
+    )