[torch_xla2] Fix the test of cumsum (#7965)

pytorch · Sep 6, 2024 · 12e5958 · 12e5958
1 parent 901c3a3
commit 12e5958
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
@@ -32,7 +32,6 @@
     "cross",
     "cummax",
     "cummin",
-    "cumsum",
     "diag",
     "diag_embed",
     "diagflat",

diff --git a/experimental/torch_xla2/torch_xla2/ops/jaten.py b/experimental/torch_xla2/torch_xla2/ops/jaten.py
@@ -556,6 +556,8 @@ def _aten_ne(x, y):
 def _aten_cumsum(x, y, dtype=None):
   if dtype:
     dtype = mappings.t2j_dtype(dtype)
+  if not x.shape:
+    return x
   res = jnp.cumsum(x, y, dtype)
   return res