Add bitwise_left_shift, bitwise_right_shift (#7343)

pytorch · Jun 25, 2024 · 557996a · 557996a
1 parent b505288
commit 557996a
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/experimental/torch_xla2/test/test_ops.py b/experimental/torch_xla2/test/test_ops.py
@@ -15,8 +15,6 @@
     "_segment_reduce",
     "_upsample_bilinear2d_aa",
     "bincount", # NOTE: dtype for int input torch gives float. This is weird.
-    "bitwise_left_shift",
-    "bitwise_right_shift",
     "block_diag",
     "broadcast_tensors",
     "broadcast_to",

diff --git a/experimental/torch_xla2/torch_xla2/ops/jaten.py b/experimental/torch_xla2/torch_xla2/ops/jaten.py
@@ -958,6 +958,18 @@ def _aten_bitwise_not(self):
   return ~self
 
 
+# aten.bitwise_left_shift
+@op(torch.ops.aten.bitwise_left_shift)
+def _aten_bitwise_left_shift(input, other):
+  return jnp.left_shift(input, other)
+
+
+# aten.bitwise_right_shift
+@op(torch.ops.aten.bitwise_right_shift)
+def _aten_bitwise_right_shift(input, other):
+  return jnp.right_shift(input, other)
+
+
 # aten.embedding_dense_backward
 
 
@@ -2113,4 +2125,4 @@ def _aten_randint(
 
 @op(torch.ops.aten.dim, is_jax_function=False)
 def _aten_dim(self):
-  return len(self.shape)
+  return len(self.shape)