Fix torch.linalg.matrix_power (#7495) (#8351)

Co-authored-by: Simon Teo <[email protected]>
pytorch · Nov 4, 2024 · 7dda437 · 7dda437
1 parent a19a996
commit 7dda437
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/experimental/torch_xla2/test/test_ops.py b/experimental/torch_xla2/test/test_ops.py
@@ -28,7 +28,6 @@
     "linalg.det",
     "linalg.ldl_solve",
     "linalg.lu_solve",
-    "linalg.matrix_power",
     "max_pool2d_with_indices_backward",
     "nn.functional.adaptive_avg_pool3d",
     "nn.functional.adaptive_max_pool1d",

diff --git a/experimental/torch_xla2/torch_xla2/ops/jtorch.py b/experimental/torch_xla2/torch_xla2/ops/jtorch.py
@@ -358,6 +358,10 @@ def linalg_solve_ex(a, b):
 def linalg_svd(a, full_matrices=True, **kwargs):
   return jaten._aten__linalg_svd(a, full_matrices=full_matrices, **kwargs)
 
+@register_function(torch.linalg.matrix_power)
+def matrix_power(A, n, *, out=None):
+  return jnp.linalg.matrix_power(A, n)
+
 @register_function(torch.svd)
 def svd(a, some=True, compute_uv=True):
   if not compute_uv: