[torch_xla2] Fix cholesky_inverse (#8214)

pytorch · Oct 4, 2024 · eba28f1 · eba28f1
1 parent 922805e
commit eba28f1
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/experimental/torch_xla2/test/test_ops.py b/experimental/torch_xla2/test/test_ops.py
@@ -21,7 +21,6 @@
  "cdist",
  "ceil",
  "cholesky",
- "cholesky_inverse",
  "cholesky_solve",
  "complex",
  "diagonal_copy",

diff --git a/experimental/torch_xla2/torch_xla2/ops/jaten.py b/experimental/torch_xla2/torch_xla2/ops/jaten.py
@@ -2041,6 +2041,14 @@ def _aten__pdist_forward(x, p=2):
  return condensed_dists
 
 
+@op(torch.ops.aten.cholesky_inverse)
+def _aten_cholesky_inverse(input, upper=False):
+ t = jnp.matrix_transpose(input)
+ if "complex" in str(input.dtype):
+ t = t.conjugate()
+ return jnp.linalg.inv(input @ t)
+
+
 # aten.cos
 @op(torch.ops.aten.cos)
 @op_base.promote_int_input