Upgrade torch dependency to 2.4.0 plus misc changes:

experimental/torch_xla2/README.md

@@ -199,7 +199,7 @@ def model_func(param, inputs):
 Now, we can apply `jax_jit`
 
 ```python
-from torch_xla2.extra import jax_jit
+from torch_xla2.interop import jax_jit
 model_func_jitted = jax_jit(model_func)
 print(model_func_jitted(new_state_dict, inputs))
 ```

experimental/torch_xla2/dev-requirements.txt

@@ -1,3 +1,3 @@
 -f https://download.pytorch.org/whl/torch
-torch==2.3.0+cpu
+torch==2.4.0+cpu
 ruff~=0.3.5

experimental/torch_xla2/test-requirements.txt

@@ -1,4 +1,6 @@
 -r dev-requirements.txt
+absl-py
+immutabledict
 pytest
 pytest-xdist
 sentencepiece

experimental/torch_xla2/test/test_exports.py

@@ -4,6 +4,7 @@
 import jax
 import jax.export
 import torch_xla2
+import torch_xla2.export
 from torch_xla2 import tensor
 from torch_xla2.ops import mappings
 
@@ -107,7 +108,8 @@ def test_export_dtypes(self):
       torch.long          : "i64",
       # NO_MAPPING        : "ui4"
       torch.uint8         : "ui8",
-      torch.uint16        : "ui16",
+      # NOTE(qihqi): torch export for uint16 seems broken at torch 2.4
+      # torch.uint16        : "ui16",
       torch.uint32        : "ui32",
       torch.uint64        : "ui64",
       # NO_MAPPING        : "f8E4M3B11FNUZ"
@@ -127,7 +129,7 @@ def test_export_dtypes(self):
     }
 
     model = TensorConstant()
-    for torch_dtype in mappings.TORCH_DTYPE_TO_JAX.keys():
+    for torch_dtype in DTYPE_TO_MLIR_STR.keys():
       if torch_dtype == None:
         ## TODO: Figure out what the None mapping should be, seems like:
         ##   torch.tensor(dtype=None) maps to f32

experimental/torch_xla2/test/test_functions.py

@@ -34,6 +34,19 @@ def test_dont_capture_conversion(self):
       t2 = self.env.to_xla(t)
       # assert no exceptions
 
+  def test_brackets(self):
+    with self.env:
+      a = torch.randn((2,3))
+      a[1] = 9
+      self.assertEqual(a[1, 0].item(), 9)
+
+  def test_bernoulli_inplace(self):
+    with self.env:
+      a = torch.randn((2,3))
+      a.bernoulli_(0.4)
+
+
+
 
 if __name__ == '__main__':
   absltest.main()

experimental/torch_xla2/torch_xla2/decompositions.py

@@ -103,6 +103,16 @@ def rand_like(self, **kwargs):
     dtype = kwargs.get('dtype')
     return torch.rand(self.shape, dtype=dtype)
 
+def channel_shuffle(self, groups):
+    batchsize, channels, height, width = self.shape
+    channels_per_group = channels // groups
+    self = self.reshape(batchsize, groups, channels_per_group, height, width)
+    self = self.transpose(1, 2)
+    self = self.reshape(batchsize, channels, height, width)
+    return self
+
+_try_register(aten.channel_shuffle, channel_shuffle)
+
 _try_register(aten.bernoulli, bernoulli)
 _try_register(aten.rand_like, rand_like)
 
@@ -121,4 +131,6 @@ def rand_like(self, **kwargs):
     torch.ops.aten.replication_pad3d,
     torch.ops.aten.bernoulli,
     torch.ops.aten.rand_like,
+    torch.ops.aten._batch_norm_with_update,
+    torch.ops.aten.channel_shuffle,
 ])

experimental/torch_xla2/torch_xla2/ops/jaten.py

@@ -23,6 +23,7 @@
   torch.ops.aten.add_: torch.ops.aten.add,
   torch.ops.aten.sub_: torch.ops.aten.sub,
   torch.ops.aten.mul_: torch.ops.aten.mul,
+  torch.ops.aten.div_: torch.ops.aten.div,
   torch.ops.aten.pow_: torch.ops.aten.pow,
   torch.ops.aten.lt_: torch.ops.aten.lt,
   torch.ops.aten.le_: torch.ops.aten.le,
@@ -34,6 +35,7 @@
   torch.ops.aten.relu_: torch.ops.aten.relu,
   torch.ops.aten.normal_: torch.ops.aten.normal,
   torch.ops.aten.squeeze_: torch.ops.aten.squeeze,
+  torch.ops.aten.bernoulli_: torch.ops.aten.bernoulli.p,
 }
 
 
@@ -277,12 +279,6 @@ def _aten_div(x, y, rounding_mode=""):
   return res
 
 
-@op(torch.ops.aten.div_, is_jax_function=False)
-def _aten_div_(x, y, rounding_mode=""):
-  x._elem = _aten_div(x._elem, y._elem, rounding_mode)
-  return x
-
-
 @op(torch.ops.aten.true_divide)
 def _aten_true_divide(x, y):
   return x / y
@@ -2219,6 +2215,19 @@ def _randn(
     res = res.astype(dtype)
   return res
 
+@op(torch.ops.aten.bernoulli.p, needs_env=True)
+def _bernoulli(
+  self,
+  p = 0.5,
+  *,
+  generator=None,
+  env=None,
+):
+  key = env.get_and_rotate_prng_key(generator)
+  res = jax.random.uniform(key, self.shape) < p
+  return res
+
+
 
 @op(torch.ops.aten.randn_like, needs_env=True)
 @op_base.convert_dtype()

experimental/torch_xla2/torch_xla2/tensor.py

@@ -104,8 +104,8 @@ def flatten(self, start_dim=0, end_dim=-1):
     # return torch.reshape(self, new_shape)
 
   def __setitem__(self, key, val):
-    key = unwrap(key)
-    self._elem = self._elem.at[key].set(val._elem)
+    key, val = self._env.t2j_iso((key, val))
+    self._elem = self._elem.at[key].set(val)
 
   def type_as(self, other):
     self._elem = self._elem.astype(other._elem.dtype)