torchax 0.0.11.dev202612__py3-none-any.whl → 0.0.11.dev2025128__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.dev202612"
+__version__ = "0.0.11.dev2025128"
 VERSION = __version__
 
 # the "fast path" uses some sparse tensor thingies that currently we

torchax/config.py

@@ -35,6 +35,10 @@ 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,6 +22,7 @@ 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
@@ -29,11 +30,6 @@ 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,16 +5776,7 @@ 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,6 +24,9 @@ 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
@@ -151,6 +154,39 @@ 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:
@@ -229,6 +265,11 @@ 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.dev202612.dist-info → torchax-0.0.11.dev2025128.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchax
-Version: 0.0.11.dev202612
+Version: 0.0.11.dev2025128
 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.dev202612.dist-info → torchax-0.0.11.dev2025128.dist-info}/RECORD

@@ -1,13 +1,13 @@
 torchax/CONTRIBUTING.md,sha256=gbB2ewxDLC-HHRBC3B8HdppV_d9MbDd-9rvzGQt7vZU,1440
-torchax/__init__.py,sha256=m9mLdO6-n-WfpeR-R7jNQxkESUUpXQJ5A-0CNnXuSJs,3889
+torchax/__init__.py,sha256=KUKGHnTP6TQi9GqMlCVo6YhIRxeryRTZA9lKyLr7vb0,3890
 torchax/amp.py,sha256=WTrfah2NYodapDVMsht7H3zDpl-XslujxhcYXr85g1s,10999
 torchax/checkpoint.py,sha256=2eoGeIQtL1Chof0W9qorB2Q0eCyEVJyWKqGoetf32GQ,2439
-torchax/config.py,sha256=oTwgWDujF9vNSHRNKUvz3ZkocDe0aDaF-yviqAJAewY,1398
+torchax/config.py,sha256=c2JVtKx-GkkQn9vGxYgweGrm57G60mH1BB-SfE6-d6Q,1497
 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=6R9_pIkd5Kwb5EoACUsq4GEEe8PaSFOdfnD4GIy6Y9U,11815
+torchax/interop.py,sha256=_cdG_lrvRD8KFkzJe5bsHqnapW2sxSABL6K2GskV-r8,11722
 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=x5vrEiVRFaTigPYvFHtR9ClwW0Z6TEHT66-mEdzfe3k,163435
+torchax/ops/jaten.py,sha256=XjYPZLX9w9fi0TtiHYqvAZQkjnjuVSNcl0ge_IYjhl8,162999
 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=OuFTed92B3WuWnfnjBBvMgwkl43PLJJiDWSMb0ZS0sw,16406
+torchax/ops/jtorch.py,sha256=WtBaDRFgw1PcXk8tW5SFOLbYNWLFULhdIQIrJv6Cy7k,17874
 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.dev202612.dist-info/METADATA,sha256=tVMhNPH26W_cwlb9oA8xPNByKjM25aY7OI4SGnGFTg4,22451
-torchax-0.0.11.dev202612.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-torchax-0.0.11.dev202612.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-torchax-0.0.11.dev202612.dist-info/RECORD,,
+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,,