From 5d446c5cf010600b7b933629a74591a6b7107944 Mon Sep 17 00:00:00 2001
From: Manfei <41607353+ManfeiBai@users.noreply.github.com>
Date: Wed, 13 Nov 2024 17:44:54 -0800
Subject: [PATCH] Update jaten.py

---
 experimental/torch_xla2/torch_xla2/ops/jaten.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/experimental/torch_xla2/torch_xla2/ops/jaten.py b/experimental/torch_xla2/torch_xla2/ops/jaten.py
index e3ed7539d31..61d15f6c2f5 100644
--- 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):