Implement torch diagflat, and movedim as well as it is also needed as…

… a helper function (#7378) (#8038)
pytorch · Sep 18, 2024 · b79a46b · b79a46b
1 parent 6b33f8f
commit b79a46b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/experimental/torch_xla2/test/test_ops.py b/experimental/torch_xla2/test/test_ops.py
@@ -26,7 +26,6 @@
     "combinations",
     "complex",
     "diag_embed",
-    "diagflat",
     "diagonal_copy",
     "diagonal_scatter",
     "digamma",

diff --git a/experimental/torch_xla2/torch_xla2/ops/jaten.py b/experimental/torch_xla2/torch_xla2/ops/jaten.py
@@ -1939,6 +1939,17 @@ def _aten_diagonal(input, offset=0, dim1=0, dim2=1):
   return jnp.diagonal(input, offset, dim1, dim2)
 
 
+# aten.diagflat
+@op(torch.ops.aten.diagflat)
+def _aten_diagflat(input, offset=0):
+  return jnp.diagflat(jnp.array(input), offset)
+
+
+@op(torch.ops.aten.movedim)
+def _aten_movedim(input, source, destination):
+  return jnp.moveaxis(input, source, destination)
+
+
 # aten.eq
 @op(torch.ops.aten.eq)
 def _aten_eq(input1, input2):