select op support in torchxla2 (#7293)

pytorch · Jun 17, 2024 · f50e951 · f50e951
1 parent 61c794d
commit f50e951
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/experimental/torch_xla2/test/test_ops.py b/experimental/torch_xla2/test/test_ops.py
@@ -290,7 +290,6 @@
     "scatter",
     "scatter_reduce",
     "searchsorted",
-    "select",
     "select_scatter",
     "signbit",
     "softmax",

diff --git a/experimental/torch_xla2/torch_xla2/ops/jaten.py b/experimental/torch_xla2/torch_xla2/ops/jaten.py
@@ -103,6 +103,9 @@ def _aten_index_copy(x, dim, indexes, source):
 
 
 @op(torch.ops.aten.select)
+def _aten_select(x, dim, indexes):
+  return jax.lax.index_in_dim(x, index=indexes, axis=dim, keepdims=False)
+
 @op(torch.ops.aten.index_select)
 @op(torch.ops.aten.select_copy)
 def _aten_index_select(x, dim, indexes):