Update jaten.py

pytorch · Nov 14, 2024 · 5d446c5 · 5d446c5
1 parent a3bb50f
commit 5d446c5
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/experimental/torch_xla2/torch_xla2/ops/jaten.py b/experimental/torch_xla2/torch_xla2/ops/jaten.py
@@ -236,6 +236,11 @@ def _aten_index_select(x, dim, index):
   return jnp.take(x, index, dim)
 
 
+@op(torch.ops.aten.cholesky)
+def _aten_cholesky(input, upper=False):
+  return jax.scipy.linalg.cholesky(input, lower=(not upper))
+
+
 # aten.igammac
 @op(torch.ops.aten.igammac)
 def _aten_igammac(input, other):