Update torch compat version to 2.7.1 (#9455)

Author: qihqi

torchax/dev-requirements.txt

@@ -1,5 +1,5 @@
 -f https://download.pytorch.org/whl/torch
-torch==2.6.0 ; sys_platform == 'darwin'  # macOS
-torch==2.6.0+cpu; sys_platform != 'darwin' # Non-macOS (CPU-only), like on TPU
+torch==2.7.1 ; sys_platform == 'darwin'  # macOS
+torch==2.7.1+cpu; sys_platform != 'darwin' # Non-macOS (CPU-only), like on TPU
 yapf==0.40.2  # N.B.: keep in sync with `infra/ansible/config/pip.yaml`, `.github/workflows/lintercheck.yml`
 flax==0.10.6

torchax/pyproject.toml

@@ -12,7 +12,7 @@ authors = [
     {name = "Han Qi", email = "[email protected]"},
     {name = "Pytorch/XLA team", email = "[email protected]"},
 ]
-description = "torchax is a library for running PyTorch on TPU"
+description = "torchax is a library for running Jax and PyTorch together"
 readme = "README.md"
 classifiers = [
     "Development Status :: 3 - Alpha",

torchax/torchax/decompositions.py

@@ -761,6 +761,13 @@ def get_summand(ix: torch.Tensor, iy: torch.Tensor, iz: torch.Tensor,
     torch.ops.aten._chunk_cat.out,
     torch.ops.aten._weight_norm_interface.default,
     torch.ops.aten._weight_norm_interface.out,
+    torch.ops.aten.__iand__.Tensor,
+    torch.ops.aten.__ixor__.Tensor,
+    torch.ops.aten.__ilshift__.Tensor,
+    torch.ops.aten.__ilshift__.Scalar,
+    torch.ops.aten.__irshift__.Tensor,
+    torch.ops.aten.__irshift__.Scalar,
+    torch.ops.aten.__ior__.Tensor,
 ])
 
 MUTABLE_DECOMPOSITION = [

torchax/torchax/ops/jaten.py

@@ -375,6 +375,7 @@ def _aten_slice(self, dim=0, start=None, end=None, step=1):
   return self[tuple(dims)]
 
 
+@op(torch.ops.aten.positive)
 @op(torch.ops.aten.detach)
 def _aten_detach(self):
   return self
@@ -2951,12 +2952,14 @@ def _aten_log2(x):
 
 # aten.logical_and
 @op(torch.ops.aten.logical_and)
+@op(torch.ops.aten.__and__)
 def _aten_logical_and(self, other):
   return jnp.logical_and(self, other)
 
 
 # aten.logical_or
 @op(torch.ops.aten.logical_or)
+@op(torch.ops.aten.__or__)
 def _aten_logical_or(self, other):
   return jnp.logical_or(self, other)
 
@@ -2998,6 +3001,7 @@ def _aten_logcumsumexp(self, dim=None):
 # aten.max_pool3d_backward
 # aten.logical_xor
 @op(torch.ops.aten.logical_xor)
+@op(torch.ops.aten.__xor__)
 def _aten_logical_xor(self, other):
   return jnp.logical_xor(self, other)
 
@@ -4946,7 +4950,7 @@ def _aten__linalg_solve_ex(a, b):
   res = jnp.linalg.solve(a, b)
   if batched:
     res = res.squeeze(-1)
-  info_shape = a.shape[0] if len(a.shape) >= 3 else []
+  info_shape = a.shape[:-2]
   info = jnp.zeros(info_shape, dtype=mappings.t2j_dtype(torch.int32))
   return res, info
 
@@ -5497,6 +5501,50 @@ def _aten_pad(self, pad, mode='constant', value=None):
     )
 
 
+@op(torch.ops.aten.is_nonzero)
+def _aten_is_nonzero(a):
+  a = jnp.squeeze(a)
+  if a.shape == (0,):
+    raise RuntimeError('bool value of Tensor with no values is ambiguous')
+  if a.ndim != 0:
+    raise RuntimeError(
+        'bool value of Tensor with more than one value is ambiguous')
+  return a.item() != 0
+
+
+@op(torch.ops.aten.logit)
+def _aten_logit(self: jax.Array, eps: float | None = None) -> jax.Array:
+  """
+  Computes the logit function of the input tensor.
+
+  logit(p) = log(p / (1 - p))
+
+  Args:
+    self: Input tensor.
+    eps: A small value to clip the input tensor to avoid log(0) or division by zero.
+         If None, no clipping is performed.
+
+  Returns:
+    A tensor with the logit of each element of the input.
+  """
+  if eps is not None:
+    self = jnp.clip(self, eps, 1.0 - eps)
+  res = jnp.log(self / (1.0 - self))
+  res = res.astype(mappings.t2j_dtype(torch.get_default_dtype()))
+  return res
+
+
+@op(torch.ops.aten.floor_divide)
+def _aten_floor_divide(x, y):
+  res = jnp.floor_divide(x, y)
+  return res
+
+
+@op(torch.ops.aten._assert_tensor_metadata)
+def _aten__assert_tensor_metadata(*args, **kwargs):
+  pass
+
+
 mutation_ops_to_functional = {
     torch.ops.aten.add_:
         op_base.InplaceOp(torch.ops.aten.add),
@@ -5565,6 +5613,10 @@ def _aten_pad(self, pad, mode='constant', value=None):
         op_base.InplaceOp(torch.ops.aten.scatter),
     torch.ops.aten.bitwise_or_:
         op_base.InplaceOp(torch.ops.aten.bitwise_or),
+    torch.ops.aten.floor_divide_:
+        op_base.InplaceOp(torch.ops.aten.floor_divide),
+    torch.ops.aten.remainder_:
+        op_base.InplaceOp(torch.ops.aten.remainder),
 }
 
 # Note: tuple comparisons work intuitively, e.g. `_jax_version >= (0, 4, 32)`.