Fix negative dims in scatter_index

pytorch · Sep 19, 2024 · 65c3333 · 65c3333
1 parent c0501f0
commit 65c3333
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/experimental/torch_xla2/test/test_ops.py b/experimental/torch_xla2/test/test_ops.py
@@ -148,7 +148,6 @@
     "resize_as_",
     "rot90",
     "rsub",
-    "scatter_add",
     "scatter",
     "scatter_reduce",
     "searchsorted",

diff --git a/experimental/torch_xla2/torch_xla2/ops/jaten.py b/experimental/torch_xla2/torch_xla2/ops/jaten.py
@@ -1343,6 +1343,8 @@ def _scatter_index(dim, index):
   index_shape = list(index.shape)
   input_indexes = []
   source_indexes = []
+  if dim < 0:
+    dim += len(index_shape)
   for i in range(len(index_shape)):
     source_indexes.append(slice(0, index_shape[i]))
     if i == dim: