torchax 0.0.11.dev2025128__py3-none-any.whl → 0.0.11.dev20251229__py3-none-any.whl

torchax/__init__.py

@@ -27,7 +27,7 @@ from torchax import tensor
 
 from .checkpoint import load_checkpoint, save_checkpoint
 
-__version__ = "0.0.11.dev2025128"
+__version__ = "0.0.11.dev20251229"
 VERSION = __version__
 
 # the "fast path" uses some sparse tensor thingies that currently we

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/interop.py

@@ -22,7 +22,6 @@ import jax
 import jax.numpy as jnp
 import torch
 from jax import tree_util as pytree
-from jax.experimental.shard_map import shard_map
 from torch.nn.utils import stateless as torch_stateless
 
 import torchax
@@ -30,6 +29,11 @@ from torchax import tensor, util
 from torchax.ops import mappings
 from torchax.types import JaxCallable, JaxValue, TorchCallable, TorchValue
 
+try:
+  from jax import shard_map as shard_map  # for jax since v0.8.0
+except ImportError:
+  from jax.experimental.shard_map import shard_map
+
 
 def extract_all_buffers(m: torch.nn.Module):
   buffers = {}

torchax/ops/jaten.py

@@ -5776,7 +5776,16 @@ mutation_ops_to_functional = {
   torch.ops.aten.scatter_add_: op_base.InplaceOp(torch.ops.aten.scatter_add),
   torch.ops.aten.scatter_reduce_.two: op_base.InplaceOp(torch.ops.aten.scatter_reduce),
   torch.ops.aten.scatter_: op_base.InplaceOp(torch.ops.aten.scatter),
+  torch.ops.aten.bitwise_not_: op_base.InplaceOp(torch.ops.aten.bitwise_not),
+  torch.ops.aten.bitwise_and_: op_base.InplaceOp(torch.ops.aten.bitwise_and),
   torch.ops.aten.bitwise_or_: op_base.InplaceOp(torch.ops.aten.bitwise_or),
+  torch.ops.aten.bitwise_xor_: op_base.InplaceOp(torch.ops.aten.bitwise_xor),
+  torch.ops.aten.bitwise_left_shift_: op_base.InplaceOp(
+    torch.ops.aten.bitwise_left_shift
+  ),
+  torch.ops.aten.bitwise_right_shift_: op_base.InplaceOp(
+    torch.ops.aten.bitwise_right_shift
+  ),
   torch.ops.aten.floor_divide_: op_base.InplaceOp(torch.ops.aten.floor_divide),
   torch.ops.aten.remainder_: op_base.InplaceOp(torch.ops.aten.remainder),
   torch.ops.aten.index_put_: op_base.InplaceOp(torch.ops.aten.index_put),

torchax/ops/jtorch.py

@@ -24,9 +24,6 @@ import jax.numpy as jnp
 import numpy as np
 import torch
 import torch.utils._pytree as pytree
-from jax.experimental.pallas.ops.tpu import flash_attention
-from jax.experimental.shard_map import shard_map
-from jax.sharding import PartitionSpec
 
 import torchax.tensor
 from torchax.ops import jaten, jimage, mappings, op_base
@@ -154,39 +151,6 @@ def _sdpa_reference(
   return attn_weight @ value
 
 
-def _tpu_flash_attention(query, key, value, env):
-  fsdp_partition = PartitionSpec("fsdp")
-
-  def wrap_flash_attention(query, key, value):
-    block_sizes = flash_attention.BlockSizes(
-      block_b=min(2, query.shape[0]),
-      block_q=min(512, query.shape[2]),
-      block_k_major=min(512, key.shape[2]),
-      block_k=min(512, key.shape[2]),
-      block_q_major_dkv=min(512, query.shape[2]),
-      block_k_major_dkv=min(512, key.shape[2]),
-      block_k_dkv=min(512, key.shape[2]),
-      block_q_dkv=min(512, query.shape[2]),
-      block_k_major_dq=min(512, key.shape[2]),
-      block_k_dq=min(256, key.shape[2]),
-      block_q_dq=min(1024, query.shape[2]),
-    )
-    return flash_attention.flash_attention(
-      query, key, value, causal=True, block_sizes=block_sizes
-    )
-
-  if env.config.shmap_flash_attention:
-    wrap_flash_attention = shard_map(
-      wrap_flash_attention,
-      mesh=env._mesh,
-      in_specs=(fsdp_partition, fsdp_partition, fsdp_partition),
-      out_specs=fsdp_partition,
-      check_rep=False,
-    )
-  # return flash_attn_mapped(query, key, value)
-  return wrap_flash_attention(query, key, value)
-
-
 @register_function(torch.nn.functional.one_hot)
 def one_hot(tensor, num_classes=-1):
   if num_classes == -1:
@@ -265,11 +229,6 @@ def scaled_dot_product_attention(
   enable_gqa=False,
   env=None,
 ) -> torch.Tensor:
-  if env.config.use_tpu_flash_attention:
-    jquery, jkey, jvalue = env.t2j_iso((query, key, value))
-    res = _tpu_flash_attention(jquery, jkey, jvalue, env)
-    return env.j2t_iso(res)
-
   return _sdpa_reference(
     query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa
   )

{torchax-0.0.11.dev2025128.dist-info → torchax-0.0.11.dev20251229.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchax
-Version: 0.0.11.dev2025128
+Version: 0.0.11.dev20251229
 Summary: torchax is a library for running Jax and PyTorch together
 Project-URL: Homepage, https://github.com/google/torchax
 Author-email: Han Qi <qihan.dev@gmail.com>, Google Cloud Inference Team <cmcs-inference-eng@google.com>

{torchax-0.0.11.dev2025128.dist-info → torchax-0.0.11.dev20251229.dist-info}/RECORD

@@ -1,13 +1,13 @@
 torchax/CONTRIBUTING.md,sha256=gbB2ewxDLC-HHRBC3B8HdppV_d9MbDd-9rvzGQt7vZU,1440
-torchax/__init__.py,sha256=KUKGHnTP6TQi9GqMlCVo6YhIRxeryRTZA9lKyLr7vb0,3890
+torchax/__init__.py,sha256=GYuxucLFbesjExbHu6xu1wjhh1up9ZdzVY8LSwv-EDU,3891
 torchax/amp.py,sha256=WTrfah2NYodapDVMsht7H3zDpl-XslujxhcYXr85g1s,10999
 torchax/checkpoint.py,sha256=2eoGeIQtL1Chof0W9qorB2Q0eCyEVJyWKqGoetf32GQ,2439
-torchax/config.py,sha256=c2JVtKx-GkkQn9vGxYgweGrm57G60mH1BB-SfE6-d6Q,1497
+torchax/config.py,sha256=oTwgWDujF9vNSHRNKUvz3ZkocDe0aDaF-yviqAJAewY,1398
 torchax/decompositions.py,sha256=8VU0FfKqbP8h3S7JzHi0iWqT5E1OrwYvuhf6cjzDTlI,29303
 torchax/device_module.py,sha256=7WrLUBjMQiAilVfRwEwJrbfkizPZAC3022UO70U5uEQ,924
 torchax/export.py,sha256=CFQESWy9-ENo9Ozf08qwqUtrNpt7Y3ubX_dP99o3FdM,9395
 torchax/flax.py,sha256=Eft46Np3qPvSLmBGOltCx8KbnrGsuBeqx3Zu0tlhMpg,1807
-torchax/interop.py,sha256=_cdG_lrvRD8KFkzJe5bsHqnapW2sxSABL6K2GskV-r8,11722
+torchax/interop.py,sha256=6R9_pIkd5Kwb5EoACUsq4GEEe8PaSFOdfnD4GIy6Y9U,11815
 torchax/mesh_util.py,sha256=Y3RVKOyLVKpbseyXTYlJlUgyNavmDAsh2MQ6pXYQDUU,9719
 torchax/tensor.py,sha256=IP7JuzUfKZVrhJdA67VdXBEh4bfypwavKRzw6llsD4U,21452
 torchax/train.py,sha256=3sqIYO1Q6GN6gRGkVVoKjjUZ3xYPgh110g5XbkaUxh4,4367
@@ -15,17 +15,17 @@ torchax/types.py,sha256=NDtW1fARypg0ZHcGRVTBZKQqxwJzWtf5E5bntng91gk,981
 torchax/util.py,sha256=Oud7oaAw1SJo_v4fwEZdjuseZ_bvngAsAQ-dOEzy_20,3675
 torchax/view.py,sha256=750VYe6tmwAINNAyjN8GDvPmgaR8luvr74szwuikGts,11256
 torchax/ops/__init__.py,sha256=uc-Rod4Xlk_oJ-rdY9P8-MGTu9dsXPombscqcSys0r8,840
-torchax/ops/jaten.py,sha256=XjYPZLX9w9fi0TtiHYqvAZQkjnjuVSNcl0ge_IYjhl8,162999
+torchax/ops/jaten.py,sha256=x5vrEiVRFaTigPYvFHtR9ClwW0Z6TEHT66-mEdzfe3k,163435
 torchax/ops/jax_reimplement.py,sha256=Te8Je2ea9jX2SFV34PNPSHVP7-z_bmFykVWeqn8Tqwo,7714
 torchax/ops/jc10d.py,sha256=sO99kYDM9WRnSENHmMkH_MXWkx6HZdDvK5Z_M0C571g,1889
 torchax/ops/jimage.py,sha256=uvYW-fMaGU6-QTXAmTZ8rmHEkkwpXh2I-bu54STf-ic,3594
 torchax/ops/jlibrary.py,sha256=ESeFS5xwV1yiCIp_yKzDXihHd1wcz5eXjFkjFKsmw3w,3470
-torchax/ops/jtorch.py,sha256=WtBaDRFgw1PcXk8tW5SFOLbYNWLFULhdIQIrJv6Cy7k,17874
+torchax/ops/jtorch.py,sha256=OuFTed92B3WuWnfnjBBvMgwkl43PLJJiDWSMb0ZS0sw,16406
 torchax/ops/jtorchvision_nms.py,sha256=VNMshE3LCsIBHVVzkrNEm0kYMF89ZVIKWlQIk6pCZB0,9197
 torchax/ops/mappings.py,sha256=ViEsZaGIi37BhuLw9hx9cA9XXl35OVuaRN6Q9yZygxk,4162
 torchax/ops/op_base.py,sha256=-rQXLpkgNZ1HM3OT1XQkvAV_7Dtq019_rAXNAg97OuE,4135
 torchax/ops/ops_registry.py,sha256=sBT41LRGmUVP4ZJ9YU1DyffatOHxe-x8oXqMhCKh0y8,1836
-torchax-0.0.11.dev2025128.dist-info/METADATA,sha256=TqAfHQIUW8aGaQHO9rwwI7ODnc8d9D0DOL6fYHscUqM,22452
-torchax-0.0.11.dev2025128.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-torchax-0.0.11.dev2025128.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-torchax-0.0.11.dev2025128.dist-info/RECORD,,
+torchax-0.0.11.dev20251229.dist-info/METADATA,sha256=LhVy_90tMdwogkXLFgi2ljDuCJlMYW94whdboa_khVA,22453
+torchax-0.0.11.dev20251229.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+torchax-0.0.11.dev20251229.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+torchax-0.0.11.dev20251229.dist-info/RECORD,,