From ff7d5b96f0292a25c15fd59b5a7bfa52945ebcea Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Thu, 7 Aug 2025 10:16:13 +0530
Subject: [PATCH 01/19] Added cholesky_inverse to all backends

---
 keras/src/backend/jax/linalg.py        | 14 +++++++++++++
 keras/src/backend/numpy/linalg.py      |  6 ++++++
 keras/src/backend/openvino/linalg.py   |  5 +++++
 keras/src/backend/tensorflow/linalg.py |  6 ++++++
 keras/src/backend/torch/linalg.py      |  2 ++
 keras/src/ops/linalg.py                | 27 ++++++++++++++++++++++++++
 keras/src/ops/linalg_test.py           | 13 +++++++++++++
 7 files changed, 73 insertions(+)

diff --git a/keras/src/backend/jax/linalg.py b/keras/src/backend/jax/linalg.py
index 05a623d89101..ed2bd8707dcb 100644
--- a/keras/src/backend/jax/linalg.py
+++ b/keras/src/backend/jax/linalg.py
@@ -25,6 +25,20 @@ def cholesky(a):
         pass
     return out
 
+def cholesky_inverse(a):
+    L = cholesky(a)
+    identity = jnp.eye(a.shape[0], dtype=a.dtype)
+    L_inv = solve_triangular(L, identity, lower=True)
+    out = L_inv.T @ L_inv
+    try:
+        if jnp.any(jnp.isnan(out)):
+            raise ValueError(
+                "Cholesky inverse failed. The input might not be a valid "
+                "positive definite matrix."
+            )
+    except jax.errors.TracerBoolConversionError:
+        pass
+    return out
 
 def det(a):
     return jnp.linalg.det(a)
diff --git a/keras/src/backend/numpy/linalg.py b/keras/src/backend/numpy/linalg.py
index 30881964f7c5..55eece178ba6 100644
--- a/keras/src/backend/numpy/linalg.py
+++ b/keras/src/backend/numpy/linalg.py
@@ -9,6 +9,12 @@
 def cholesky(a):
     return np.linalg.cholesky(a)
 
+def cholesky_inverse(a):
+    L = np.linalg.cholesky(a)
+    identity = np.eye(a.shape[0], dtype=a.dtype)
+    L_inv = np.linalg.solve(L, identity)
+    a_inv = L_inv.T @ L_inv
+    return a_inv
 
 def det(a):
     return np.linalg.det(a)
diff --git a/keras/src/backend/openvino/linalg.py b/keras/src/backend/openvino/linalg.py
index 3703bd83a0c1..d0c889c80f67 100644
--- a/keras/src/backend/openvino/linalg.py
+++ b/keras/src/backend/openvino/linalg.py
@@ -3,6 +3,11 @@ def cholesky(a):
         "`cholesky` is not supported with openvino backend"
     )
 
+def cholesky_inverse(a):
+    raise NotImplementedError(
+        "`cholesky inverse` is not supported with openvino backend"
+    )
+
 
 def det(a):
     raise NotImplementedError("`det` is not supported with openvino backend")
diff --git a/keras/src/backend/tensorflow/linalg.py b/keras/src/backend/tensorflow/linalg.py
index da1ff4259685..eafeb9694cf8 100644
--- a/keras/src/backend/tensorflow/linalg.py
+++ b/keras/src/backend/tensorflow/linalg.py
@@ -12,6 +12,12 @@ def cholesky(a):
     # tf.linalg.cholesky simply returns NaNs for non-positive definite matrices
     return tf.debugging.check_numerics(out, "Cholesky")
 
+def cholesky_inverse(a):
+    L = cholesky(a)
+    identity = tf.eye(num_rows=tf.shape(a)[-1], dtype=a.dtype)
+    L_inv = solve_triangular(L, identity, lower=True)
+    a_inv = tf.matmul(L_inv, L_inv, transpose_a=True)
+    return a_inv
 
 def det(a):
     return tf.linalg.det(a)
diff --git a/keras/src/backend/torch/linalg.py b/keras/src/backend/torch/linalg.py
index 939074a680cd..2cf2b9891261 100644
--- a/keras/src/backend/torch/linalg.py
+++ b/keras/src/backend/torch/linalg.py
@@ -10,6 +10,8 @@
 def cholesky(x):
     return torch.linalg.cholesky(x)
 
+def cholesky_inverse(x):
+    return torch.cholesky_inverse(x)
 
 def det(x):
     return torch.det(x)
diff --git a/keras/src/ops/linalg.py b/keras/src/ops/linalg.py
index dc8004f309fd..6e23c6b739dd 100644
--- a/keras/src/ops/linalg.py
+++ b/keras/src/ops/linalg.py
@@ -43,6 +43,33 @@ def _cholesky(x):
         raise ValueError(f"Cholesky decomposition failed: {e}")
 
 
+class CholeskyInverse(Operation):
+    def call(self, x):
+        return _cholesky_inverse(x)
+
+    def compute_output_spec(self, x):
+        _assert_2d(x)
+        _assert_square(x)
+        return KerasTensor(x.shape, x.dtype)
+
+
+@keras_export(["keras.ops.cholesky_inverse", "keras.ops.linalg.cholesky_inverse"])
+def cholesky_inverse(x):
+    if any_symbolic_tensors((x,)):
+        return CholeskyInverse().symbolic_call(x)
+    return _cholesky_inverse(x)
+
+
+def _cholesky_inverse(x):
+    x = backend.convert_to_tensor(x)
+    _assert_2d(x)
+    _assert_square(x)
+    try:
+        return backend.linalg.cholesky_inverse(x)
+    except Exception as e:
+        raise ValueError(f"Cholesky inverse failed: {e}")
+
+
 class Det(Operation):
     def call(self, x):
         return _det(x)
diff --git a/keras/src/ops/linalg_test.py b/keras/src/ops/linalg_test.py
index 67d72b32eee8..6d9af2a38f9d 100644
--- a/keras/src/ops/linalg_test.py
+++ b/keras/src/ops/linalg_test.py
@@ -23,6 +23,19 @@ def test_cholesky(self):
         with self.assertRaises(ValueError):
             linalg.cholesky(x)
 
+    def test_cholesky_inverse(self):
+        x = KerasTensor([None, 20, 20])
+        out = linalg.cholesky_inverse(x)
+        self.assertEqual(out.shape, (None, 20, 20))
+
+        x = KerasTensor([None, None, 20])
+        with self.assertRaises(ValueError):
+            linalg.cholesky_inverse(x)
+
+        x = KerasTensor([None, 20, 15])
+        with self.assertRaises(ValueError):
+            linalg.cholesky_inverse(x)
+
     def test_det(self):
         x = KerasTensor([None, 20, 20])
         out = linalg.det(x)

From 26a15f0beee591443d72f9b184227e5a5dec7b48 Mon Sep 17 00:00:00 2001
From: Suhana <suhanaaa@google.com>
Date: Thu, 7 Aug 2025 10:46:38 +0530
Subject: [PATCH 02/19] Update keras/src/backend/openvino/linalg.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 keras/src/backend/openvino/linalg.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/keras/src/backend/openvino/linalg.py b/keras/src/backend/openvino/linalg.py
index d0c889c80f67..8ecadd33b2f9 100644
--- a/keras/src/backend/openvino/linalg.py
+++ b/keras/src/backend/openvino/linalg.py
@@ -4,9 +4,9 @@ def cholesky(a):
     )
 
 def cholesky_inverse(a):
-    raise NotImplementedError(
-        "`cholesky inverse` is not supported with openvino backend"
-    )
+raise NotImplementedError(
+    "`Cholesky inverse` is not supported with the OpenVINO backend"
+)
 
 
 def det(a):

From 6cdaf05afde990b5fdbce7a21670be1f1e2ab1a5 Mon Sep 17 00:00:00 2001
From: Suhana <suhanaaa@google.com>
Date: Thu, 7 Aug 2025 10:47:30 +0530
Subject: [PATCH 03/19] Update keras/src/ops/linalg.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 keras/src/ops/linalg.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/keras/src/ops/linalg.py b/keras/src/ops/linalg.py
index 6e23c6b739dd..9f9c5d8352b8 100644
--- a/keras/src/ops/linalg.py
+++ b/keras/src/ops/linalg.py
@@ -53,8 +53,23 @@ def compute_output_spec(self, x):
         return KerasTensor(x.shape, x.dtype)
 
 
-@keras_export(["keras.ops.cholesky_inverse", "keras.ops.linalg.cholesky_inverse"])
 def cholesky_inverse(x):
+    """Computes the inverse of a symmetric positive-definite matrix using the
+    Cholesky decomposition.
+
+    This function is more efficient and numerically stable than `keras.ops.inv`
+    for symmetric positive-definite matrices.
+
+    Args:
+        x: Input tensor of shape `(..., M, M)`. The matrix must be symmetric and
+            positive-definite.
+
+    Returns:
+        A tensor of shape `(..., M, M)` representing the inverse of `x`.
+
+    Raises:
+        ValueError: If `x` is not a symmetric positive-definite matrix.
+    """
     if any_symbolic_tensors((x,)):
         return CholeskyInverse().symbolic_call(x)
     return _cholesky_inverse(x)

From b280d12af77a8d410da71a62a88e7e95b935f0da Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Thu, 7 Aug 2025 14:59:50 +0530
Subject: [PATCH 04/19] Added more tests for cholesky inverse op

---
 keras/src/backend/jax/linalg.py        |  7 +++----
 keras/src/backend/numpy/linalg.py      | 11 +++++------
 keras/src/backend/tensorflow/linalg.py |  7 +++----
 keras/src/backend/torch/linalg.py      |  2 ++
 keras/src/ops/linalg_test.py           | 19 +++++++++++++++++++
 5 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/keras/src/backend/jax/linalg.py b/keras/src/backend/jax/linalg.py
index ed2bd8707dcb..8950e2986df2 100644
--- a/keras/src/backend/jax/linalg.py
+++ b/keras/src/backend/jax/linalg.py
@@ -26,10 +26,9 @@ def cholesky(a):
     return out
 
 def cholesky_inverse(a):
-    L = cholesky(a)
-    identity = jnp.eye(a.shape[0], dtype=a.dtype)
-    L_inv = solve_triangular(L, identity, lower=True)
-    out = L_inv.T @ L_inv
+    identity = jnp.eye(a.shape[-1], dtype=a.dtype)
+    a_inv = solve_triangular(a, identity, lower=True)
+    out = jnp.matmul(jnp.transpose(a_inv), a_inv)
     try:
         if jnp.any(jnp.isnan(out)):
             raise ValueError(
diff --git a/keras/src/backend/numpy/linalg.py b/keras/src/backend/numpy/linalg.py
index 55eece178ba6..ec1b1c1e1d55 100644
--- a/keras/src/backend/numpy/linalg.py
+++ b/keras/src/backend/numpy/linalg.py
@@ -9,12 +9,11 @@
 def cholesky(a):
     return np.linalg.cholesky(a)
 
-def cholesky_inverse(a):
-    L = np.linalg.cholesky(a)
-    identity = np.eye(a.shape[0], dtype=a.dtype)
-    L_inv = np.linalg.solve(L, identity)
-    a_inv = L_inv.T @ L_inv
-    return a_inv
+def cholesky_inverse(a):    
+    identity = np.eye(a.shape[-1], dtype=a.dtype)
+    a_inv = solve_triangular(a, identity, lower=True)
+    out = solve_triangular(np.transpose(a), a_inv, lower=False)
+    return out
 
 def det(a):
     return np.linalg.det(a)
diff --git a/keras/src/backend/tensorflow/linalg.py b/keras/src/backend/tensorflow/linalg.py
index eafeb9694cf8..f98e2db7c7a4 100644
--- a/keras/src/backend/tensorflow/linalg.py
+++ b/keras/src/backend/tensorflow/linalg.py
@@ -13,11 +13,10 @@ def cholesky(a):
     return tf.debugging.check_numerics(out, "Cholesky")
 
 def cholesky_inverse(a):
-    L = cholesky(a)
     identity = tf.eye(num_rows=tf.shape(a)[-1], dtype=a.dtype)
-    L_inv = solve_triangular(L, identity, lower=True)
-    a_inv = tf.matmul(L_inv, L_inv, transpose_a=True)
-    return a_inv
+    a_inv = solve_triangular(a, identity, lower=True)
+    out = tf.matmul(a_inv, a_inv, transpose_a=True)
+    return out
 
 def det(a):
     return tf.linalg.det(a)
diff --git a/keras/src/backend/torch/linalg.py b/keras/src/backend/torch/linalg.py
index 2cf2b9891261..54283e44cd0b 100644
--- a/keras/src/backend/torch/linalg.py
+++ b/keras/src/backend/torch/linalg.py
@@ -10,9 +10,11 @@
 def cholesky(x):
     return torch.linalg.cholesky(x)
 
+
 def cholesky_inverse(x):
     return torch.cholesky_inverse(x)
 
+
 def det(x):
     return torch.det(x)
 
diff --git a/keras/src/ops/linalg_test.py b/keras/src/ops/linalg_test.py
index 6d9af2a38f9d..7b269e08deb3 100644
--- a/keras/src/ops/linalg_test.py
+++ b/keras/src/ops/linalg_test.py
@@ -209,6 +209,15 @@ def test_cholesky(self):
         with self.assertRaises(ValueError):
             linalg.cholesky(x)
 
+    def test_cholesky_inverse(self):
+        x = KerasTensor([4, 3, 3])
+        out = linalg.cholesky_inverse(x)
+        self.assertEqual(out.shape, (4, 3, 3))
+
+        x = KerasTensor([10, 20, 15])
+        with self.assertRaises(ValueError):
+            linalg.cholesky_inverse(x)
+
     def test_det(self):
         x = KerasTensor([4, 3, 3])
         out = linalg.det(x)
@@ -351,6 +360,16 @@ def test_cholesky(self):
         out = linalg.cholesky(x_psd)
         self.assertAllClose(out, np.linalg.cholesky(x_psd), atol=1e-4)
 
+    def test_cholesky_inverse(self):
+        x_np = np.random.rand(3, 3).astype("float32")
+        x_psd_np = x_np @ x_np.T + 1e-4 * np.eye(3, dtype="float32")
+        x = linalg.cholesky(x_psd_np)
+        result_from_op = linalg.cholesky_inverse(x)
+        ground_truth_np = np.linalg.inv(x_psd_np)
+        self.assertAllClose(
+            result_from_op, ground_truth_np, atol=1e-5, rtol=1e-5
+        )
+
     def test_det(self):
         x = np.random.rand(4, 3, 3)
         out = linalg.det(x)

From 393a52b624a6af769c14b9e920de15586fb8daed Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Thu, 7 Aug 2025 21:07:33 +0530
Subject: [PATCH 05/19] Addressed the comments

---
 keras/src/backend/jax/linalg.py      | 8 --------
 keras/src/backend/openvino/linalg.py | 7 ++++---
 keras/src/ops/linalg.py              | 1 +
 3 files changed, 5 insertions(+), 11 deletions(-)

diff --git a/keras/src/backend/jax/linalg.py b/keras/src/backend/jax/linalg.py
index 8950e2986df2..981aea27368a 100644
--- a/keras/src/backend/jax/linalg.py
+++ b/keras/src/backend/jax/linalg.py
@@ -29,14 +29,6 @@ def cholesky_inverse(a):
     identity = jnp.eye(a.shape[-1], dtype=a.dtype)
     a_inv = solve_triangular(a, identity, lower=True)
     out = jnp.matmul(jnp.transpose(a_inv), a_inv)
-    try:
-        if jnp.any(jnp.isnan(out)):
-            raise ValueError(
-                "Cholesky inverse failed. The input might not be a valid "
-                "positive definite matrix."
-            )
-    except jax.errors.TracerBoolConversionError:
-        pass
     return out
 
 def det(a):
diff --git a/keras/src/backend/openvino/linalg.py b/keras/src/backend/openvino/linalg.py
index 8ecadd33b2f9..e794feac4ae0 100644
--- a/keras/src/backend/openvino/linalg.py
+++ b/keras/src/backend/openvino/linalg.py
@@ -3,10 +3,11 @@ def cholesky(a):
         "`cholesky` is not supported with openvino backend"
     )
 
+
 def cholesky_inverse(a):
-raise NotImplementedError(
-    "`Cholesky inverse` is not supported with the OpenVINO backend"
-)
+    raise NotImplementedError(
+        "`Cholesky inverse` is not supported with the OpenVINO backend."
+    )
 
 
 def det(a):
diff --git a/keras/src/ops/linalg.py b/keras/src/ops/linalg.py
index 9f9c5d8352b8..bdd83cba0e2d 100644
--- a/keras/src/ops/linalg.py
+++ b/keras/src/ops/linalg.py
@@ -53,6 +53,7 @@ def compute_output_spec(self, x):
         return KerasTensor(x.shape, x.dtype)
 
 
+@keras_export(["keras.ops.cholesky_inverse", "keras.ops.linalg.cholesky_inverse"])
 def cholesky_inverse(x):
     """Computes the inverse of a symmetric positive-definite matrix using the
     Cholesky decomposition.

From 29fe66c3526e42e14d7fa0621693b9c5c11e8b50 Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Thu, 7 Aug 2025 21:28:48 +0530
Subject: [PATCH 06/19] Formatted the code for extra spaces and modified  test
 for numerical verification

---
 keras/src/backend/jax/linalg.py        |  2 ++
 keras/src/backend/numpy/linalg.py      |  2 ++
 keras/src/backend/tensorflow/linalg.py |  2 ++
 keras/src/ops/linalg_test.py           | 11 ++++++-----
 4 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/keras/src/backend/jax/linalg.py b/keras/src/backend/jax/linalg.py
index 981aea27368a..7198ef6a6529 100644
--- a/keras/src/backend/jax/linalg.py
+++ b/keras/src/backend/jax/linalg.py
@@ -25,12 +25,14 @@ def cholesky(a):
         pass
     return out
 
+
 def cholesky_inverse(a):
     identity = jnp.eye(a.shape[-1], dtype=a.dtype)
     a_inv = solve_triangular(a, identity, lower=True)
     out = jnp.matmul(jnp.transpose(a_inv), a_inv)
     return out
 
+
 def det(a):
     return jnp.linalg.det(a)
 
diff --git a/keras/src/backend/numpy/linalg.py b/keras/src/backend/numpy/linalg.py
index ec1b1c1e1d55..89305263c7a5 100644
--- a/keras/src/backend/numpy/linalg.py
+++ b/keras/src/backend/numpy/linalg.py
@@ -9,12 +9,14 @@
 def cholesky(a):
     return np.linalg.cholesky(a)
 
+
 def cholesky_inverse(a):    
     identity = np.eye(a.shape[-1], dtype=a.dtype)
     a_inv = solve_triangular(a, identity, lower=True)
     out = solve_triangular(np.transpose(a), a_inv, lower=False)
     return out
 
+
 def det(a):
     return np.linalg.det(a)
 
diff --git a/keras/src/backend/tensorflow/linalg.py b/keras/src/backend/tensorflow/linalg.py
index f98e2db7c7a4..9e64bf0ac900 100644
--- a/keras/src/backend/tensorflow/linalg.py
+++ b/keras/src/backend/tensorflow/linalg.py
@@ -12,12 +12,14 @@ def cholesky(a):
     # tf.linalg.cholesky simply returns NaNs for non-positive definite matrices
     return tf.debugging.check_numerics(out, "Cholesky")
 
+
 def cholesky_inverse(a):
     identity = tf.eye(num_rows=tf.shape(a)[-1], dtype=a.dtype)
     a_inv = solve_triangular(a, identity, lower=True)
     out = tf.matmul(a_inv, a_inv, transpose_a=True)
     return out
 
+
 def det(a):
     return tf.linalg.det(a)
 
diff --git a/keras/src/ops/linalg_test.py b/keras/src/ops/linalg_test.py
index 7b269e08deb3..30ddb10e61dd 100644
--- a/keras/src/ops/linalg_test.py
+++ b/keras/src/ops/linalg_test.py
@@ -362,12 +362,13 @@ def test_cholesky(self):
 
     def test_cholesky_inverse(self):
         x_np = np.random.rand(3, 3).astype("float32")
-        x_psd_np = x_np @ x_np.T + 1e-4 * np.eye(3, dtype="float32")
-        x = linalg.cholesky(x_psd_np)
-        result_from_op = linalg.cholesky_inverse(x)
-        ground_truth_np = np.linalg.inv(x_psd_np)
+        x_psd_np = np.matmul(x_np, x_np.T) + 1e-4 * np.eye(3, dtype="float32")
+        x_chol = linalg.cholesky(x_psd_np)
+        x_inv = linalg.cholesky_inverse(x_chol)
+        reconstructed_identity = ops.matmul(x_psd_np, x_inv)
+        identity = np.eye(3, dtype="float32")
         self.assertAllClose(
-            result_from_op, ground_truth_np, atol=1e-5, rtol=1e-5
+            reconstructed_identity, identity, atol=1e-4, rtol=1e-4
         )
 
     def test_det(self):

From d246fbe004b0f1edc84ddbc9ca809668d0858159 Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Fri, 8 Aug 2025 09:14:33 +0530
Subject: [PATCH 07/19] Added upper bool to the cholesky and cholesky_inverse
 operator

---
 keras/src/backend/jax/linalg.py        | 16 ++++++---
 keras/src/backend/numpy/linalg.py      | 16 ++++++---
 keras/src/backend/openvino/linalg.py   |  6 ++--
 keras/src/backend/tensorflow/linalg.py | 19 ++++++----
 keras/src/backend/torch/linalg.py      |  8 +++--
 keras/src/ops/linalg.py                | 49 ++++++++++++--------------
 keras/src/ops/linalg_test.py           | 34 +++++++++++++-----
 7 files changed, 92 insertions(+), 56 deletions(-)

diff --git a/keras/src/backend/jax/linalg.py b/keras/src/backend/jax/linalg.py
index 7198ef6a6529..acd992d5af44 100644
--- a/keras/src/backend/jax/linalg.py
+++ b/keras/src/backend/jax/linalg.py
@@ -9,7 +9,7 @@
 from keras.src.backend.jax.core import convert_to_tensor
 
 
-def cholesky(a):
+def cholesky(a, upper=False):
     out = jnp.linalg.cholesky(a)
     try:
         # In eager mode, raise for nan to
@@ -23,14 +23,20 @@ def cholesky(a):
     except jax.errors.TracerBoolConversionError:
         # Cannot raise for nan in tracing mode
         pass
+    if upper:
+        return jnp.swapaxes(out, -2, -1)
     return out
 
 
-def cholesky_inverse(a):
+def cholesky_inverse(a, upper=False):
     identity = jnp.eye(a.shape[-1], dtype=a.dtype)
-    a_inv = solve_triangular(a, identity, lower=True)
-    out = jnp.matmul(jnp.transpose(a_inv), a_inv)
-    return out
+    if upper:
+        u_inv = solve_triangular(a, identity, lower=False)
+        a_inv = jnp.matmul(u_inv, jnp.transpose(u_inv))
+    else:
+        l_inv = solve_triangular(a, identity, lower=True)
+        a_inv = jnp.matmul(jnp.transpose(l_inv), l_inv)
+    return a_inv
 
 
 def det(a):
diff --git a/keras/src/backend/numpy/linalg.py b/keras/src/backend/numpy/linalg.py
index 89305263c7a5..ea4bc16dac2e 100644
--- a/keras/src/backend/numpy/linalg.py
+++ b/keras/src/backend/numpy/linalg.py
@@ -6,15 +6,21 @@
 from keras.src.backend.numpy.core import convert_to_tensor
 
 
-def cholesky(a):
+def cholesky(a, upper=False):
+    if upper:
+        return np.linalg.cholesky(a, upper=True)
     return np.linalg.cholesky(a)
 
 
-def cholesky_inverse(a):    
+def cholesky_inverse(a, upper=False):
     identity = np.eye(a.shape[-1], dtype=a.dtype)
-    a_inv = solve_triangular(a, identity, lower=True)
-    out = solve_triangular(np.transpose(a), a_inv, lower=False)
-    return out
+    if upper:
+        u_inv = solve_triangular(a, identity, lower=False)
+        a_inv = np.matmul(u_inv, u_inv.T)
+    else:
+        l_inv = solve_triangular(a, identity, lower=True)
+        a_inv = np.matmul(l_inv.T, l_inv)
+    return a_inv
 
 
 def det(a):
diff --git a/keras/src/backend/openvino/linalg.py b/keras/src/backend/openvino/linalg.py
index e794feac4ae0..f4dae1c17b5f 100644
--- a/keras/src/backend/openvino/linalg.py
+++ b/keras/src/backend/openvino/linalg.py
@@ -1,12 +1,12 @@
-def cholesky(a):
+def cholesky(a, upper=False):
     raise NotImplementedError(
         "`cholesky` is not supported with openvino backend"
     )
 
 
-def cholesky_inverse(a):
+def cholesky_inverse(a, upper=False):
     raise NotImplementedError(
-        "`Cholesky inverse` is not supported with the OpenVINO backend."
+        "`Cholesky inverse` is not supported with the openvino backend"
     )
 
 
diff --git a/keras/src/backend/tensorflow/linalg.py b/keras/src/backend/tensorflow/linalg.py
index 9e64bf0ac900..eb22f3c3f480 100644
--- a/keras/src/backend/tensorflow/linalg.py
+++ b/keras/src/backend/tensorflow/linalg.py
@@ -7,17 +7,24 @@
 from keras.src.backend.tensorflow.core import convert_to_tensor
 
 
-def cholesky(a):
+def cholesky(a, upper=True):
     out = tf.linalg.cholesky(a)
     # tf.linalg.cholesky simply returns NaNs for non-positive definite matrices
-    return tf.debugging.check_numerics(out, "Cholesky")
+    out = tf.debugging.check_numerics(out, "Cholesky")
+    if upper:
+        return tf.linalg.matrix_transpose(out)
+    return out
 
 
-def cholesky_inverse(a):
+def cholesky_inverse(a, upper=True):
     identity = tf.eye(num_rows=tf.shape(a)[-1], dtype=a.dtype)
-    a_inv = solve_triangular(a, identity, lower=True)
-    out = tf.matmul(a_inv, a_inv, transpose_a=True)
-    return out
+    if upper:
+        u_inv = tf.linalg.triangular_solve(a, identity, lower=False)
+        a_inv = tf.matmul(u_inv, u_inv, transpose_b=True)
+    else:
+        l_inv = tf.linalg.triangular_solve(a, identity, lower=True)
+        a_inv = tf.matmul(l_inv, l_inv, transpose_a=True)
+    return a_inv
 
 
 def det(a):
diff --git a/keras/src/backend/torch/linalg.py b/keras/src/backend/torch/linalg.py
index 54283e44cd0b..2ccbc1c3e4e3 100644
--- a/keras/src/backend/torch/linalg.py
+++ b/keras/src/backend/torch/linalg.py
@@ -7,11 +7,15 @@
 from keras.src.backend.torch.core import convert_to_tensor
 
 
-def cholesky(x):
+def cholesky(x, upper=False):
+    if upper:
+        return torch.linalg.cholesky(x, upper=True)
     return torch.linalg.cholesky(x)
 
 
-def cholesky_inverse(x):
+def cholesky_inverse(x, upper=False):
+    if upper:
+        return torch.cholesky_inverse(x, upper=True)
     return torch.cholesky_inverse(x)
 
 
diff --git a/keras/src/ops/linalg.py b/keras/src/ops/linalg.py
index bdd83cba0e2d..8dd9ff0b73f6 100644
--- a/keras/src/ops/linalg.py
+++ b/keras/src/ops/linalg.py
@@ -7,63 +7,60 @@
 
 
 class Cholesky(Operation):
-    def call(self, x):
-        return _cholesky(x)
+    def call(self, x, upper=False):
+        return _cholesky(x, upper)
 
-    def compute_output_spec(self, x):
+    def compute_output_spec(self, x, upper=False):
         _assert_2d(x)
         _assert_square(x)
         return KerasTensor(x.shape, x.dtype)
 
 
 @keras_export(["keras.ops.cholesky", "keras.ops.linalg.cholesky"])
-def cholesky(x):
+def cholesky(x, upper=False):
     """Computes the Cholesky decomposition of a positive semi-definite matrix.
 
     Args:
         x: Input tensor of shape `(..., M, M)`.
+        upper (bool): If True, returns the upper-triangular Cholesky factor.
+            If False (default), returns the lower-triangular Cholesky factor.
 
     Returns:
-        A tensor of shape `(..., M, M)` representing the lower triangular
-        Cholesky factor of `x`.
-
+        A tensor of shape `(..., M, M)` representing the Cholesky factor of `x`.
     """
     if any_symbolic_tensors((x,)):
-        return Cholesky().symbolic_call(x)
-    return _cholesky(x)
+        return Cholesky().symbolic_call(x, upper=upper)
+    return _cholesky(x, upper=upper)
 
 
-def _cholesky(x):
+def _cholesky(x, upper=False):
     x = backend.convert_to_tensor(x)
     _assert_2d(x)
     _assert_square(x)
     try:
-        return backend.linalg.cholesky(x)
+        return backend.linalg.cholesky(x, upper=upper)
     except Exception as e:
         raise ValueError(f"Cholesky decomposition failed: {e}")
 
 
 class CholeskyInverse(Operation):
-    def call(self, x):
-        return _cholesky_inverse(x)
+    def call(self, x, upper=False):
+        return _cholesky_inverse(x, upper)
 
-    def compute_output_spec(self, x):
+    def compute_output_spec(self, x, upper=False):
         _assert_2d(x)
         _assert_square(x)
         return KerasTensor(x.shape, x.dtype)
 
 
 @keras_export(["keras.ops.cholesky_inverse", "keras.ops.linalg.cholesky_inverse"])
-def cholesky_inverse(x):
-    """Computes the inverse of a symmetric positive-definite matrix using the
-    Cholesky decomposition.
-
-    This function is more efficient and numerically stable than `keras.ops.inv`
-    for symmetric positive-definite matrices.
+def cholesky_inverse(x, upper=False):
+    """Computes the inverse of a symmetric positive-definite matrix.
 
     Args:
-        x: Input tensor of shape `(..., M, M)`. The matrix must be symmetric and
-            positive-definite.
+        x: Input tensor of shape `(..., M, M)`.
+        upper (bool): Determines whether to use the upper- or lower-triangular
+            factor for the internal computation. Defaults to False.
 
     Returns:
         A tensor of shape `(..., M, M)` representing the inverse of `x`.
@@ -72,16 +69,16 @@ def cholesky_inverse(x):
         ValueError: If `x` is not a symmetric positive-definite matrix.
     """
     if any_symbolic_tensors((x,)):
-        return CholeskyInverse().symbolic_call(x)
-    return _cholesky_inverse(x)
+        return CholeskyInverse().symbolic_call(x, upper=upper)
+    return _cholesky_inverse(x, upper=upper)
 
 
-def _cholesky_inverse(x):
+def _cholesky_inverse(x, upper=False):
     x = backend.convert_to_tensor(x)
     _assert_2d(x)
     _assert_square(x)
     try:
-        return backend.linalg.cholesky_inverse(x)
+        return backend.linalg.cholesky_inverse(x, upper=upper)
     except Exception as e:
         raise ValueError(f"Cholesky inverse failed: {e}")
 
diff --git a/keras/src/ops/linalg_test.py b/keras/src/ops/linalg_test.py
index 30ddb10e61dd..edbf80876dd6 100644
--- a/keras/src/ops/linalg_test.py
+++ b/keras/src/ops/linalg_test.py
@@ -353,22 +353,38 @@ def test_svd(self):
 
 class LinalgOpsCorrectnessTest(testing.TestCase):
     def test_cholesky(self):
-        x = np.random.rand(4, 3, 3).astype("float32")
+        x_non_psd = np.random.rand(4, 3, 3).astype("float32")
         with self.assertRaises(ValueError):
-            linalg.cholesky(x)
-        x_psd = x @ x.transpose((0, 2, 1)) + 1e-5 * np.eye(3)
-        out = linalg.cholesky(x_psd)
-        self.assertAllClose(out, np.linalg.cholesky(x_psd), atol=1e-4)
+            linalg.cholesky(x_non_psd)
+        x = np.random.rand(4, 3, 3).astype("float32")
+        x_psd = np.matmul(x, x.transpose((0, 2, 1))) + 1e-5 * np.eye(3)
+
+        l_out = linalg.cholesky(x_psd, upper=False)
+        l_expected = np.linalg.cholesky(x_psd)
+        self.assertAllClose(l_out, l_expected, atol=1e-4)
+
+        u_out = linalg.cholesky(x_psd, upper=True)
+        u_expected = l_expected.transpose((0, 2, 1))
+        self.assertAllClose(u_out, u_expected, atol=1e-4)
+
 
     def test_cholesky_inverse(self):
         x_np = np.random.rand(3, 3).astype("float32")
         x_psd_np = np.matmul(x_np, x_np.T) + 1e-4 * np.eye(3, dtype="float32")
-        x_chol = linalg.cholesky(x_psd_np)
-        x_inv = linalg.cholesky_inverse(x_chol)
-        reconstructed_identity = ops.matmul(x_psd_np, x_inv)
         identity = np.eye(3, dtype="float32")
+
+        l_factor = linalg.cholesky(x_psd_np, upper=False)
+        x_inv_from_l = linalg.cholesky_inverse(l_factor, upper=False)
+        reconstructed_from_l = ops.matmul(x_psd_np, x_inv_from_l)
+        self.assertAllClose(
+            reconstructed_from_l, identity, atol=1e-4, rtol=1e-4
+        )
+
+        u_factor = linalg.cholesky(x_psd_np, upper=True)
+        x_inv_from_u = linalg.cholesky_inverse(u_factor, upper=True)
+        reconstructed_from_u = ops.matmul(x_psd_np, x_inv_from_u)
         self.assertAllClose(
-            reconstructed_identity, identity, atol=1e-4, rtol=1e-4
+            reconstructed_from_u, identity, atol=1e-4, rtol=1e-4
         )
 
     def test_det(self):

From 30a91dda71c50498b9b8dcb8ff78aa450f178e6c Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Fri, 8 Aug 2025 10:18:03 +0530
Subject: [PATCH 08/19] Modified test case for numerical verification

---
 keras/src/backend/openvino/linalg.py | 2 +-
 keras/src/ops/linalg_test.py         | 8 ++++----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/keras/src/backend/openvino/linalg.py b/keras/src/backend/openvino/linalg.py
index f4dae1c17b5f..eda220013f58 100644
--- a/keras/src/backend/openvino/linalg.py
+++ b/keras/src/backend/openvino/linalg.py
@@ -6,7 +6,7 @@ def cholesky(a, upper=False):
 
 def cholesky_inverse(a, upper=False):
     raise NotImplementedError(
-        "`Cholesky inverse` is not supported with the openvino backend"
+        "`cholesky_inverse` is not supported with the openvino backend"
     )
 
 
diff --git a/keras/src/ops/linalg_test.py b/keras/src/ops/linalg_test.py
index edbf80876dd6..2f7577161a27 100644
--- a/keras/src/ops/linalg_test.py
+++ b/keras/src/ops/linalg_test.py
@@ -373,15 +373,15 @@ def test_cholesky_inverse(self):
         x_psd_np = np.matmul(x_np, x_np.T) + 1e-4 * np.eye(3, dtype="float32")
         identity = np.eye(3, dtype="float32")
 
-        l_factor = linalg.cholesky(x_psd_np, upper=False)
-        x_inv_from_l = linalg.cholesky_inverse(l_factor, upper=False)
+        l_factor_np = np.linalg.cholesky(x_psd_np)
+        x_inv_from_l = linalg.cholesky_inverse(l_factor_np, upper=False)
         reconstructed_from_l = ops.matmul(x_psd_np, x_inv_from_l)
         self.assertAllClose(
             reconstructed_from_l, identity, atol=1e-4, rtol=1e-4
         )
 
-        u_factor = linalg.cholesky(x_psd_np, upper=True)
-        x_inv_from_u = linalg.cholesky_inverse(u_factor, upper=True)
+        u_factor_np = l_factor_np.T
+        x_inv_from_u = linalg.cholesky_inverse(u_factor_np, upper=True)
         reconstructed_from_u = ops.matmul(x_psd_np, x_inv_from_u)
         self.assertAllClose(
             reconstructed_from_u, identity, atol=1e-4, rtol=1e-4

From 82c58b2d79f170eea0e411b3401bf0d2d9bedb43 Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Fri, 8 Aug 2025 10:20:25 +0530
Subject: [PATCH 09/19] nit changes in the openvino linalg file

---
 keras/src/backend/openvino/linalg.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/keras/src/backend/openvino/linalg.py b/keras/src/backend/openvino/linalg.py
index eda220013f58..31ae0abe37c6 100644
--- a/keras/src/backend/openvino/linalg.py
+++ b/keras/src/backend/openvino/linalg.py
@@ -6,7 +6,7 @@ def cholesky(a, upper=False):
 
 def cholesky_inverse(a, upper=False):
     raise NotImplementedError(
-        "`cholesky_inverse` is not supported with the openvino backend"
+        "`cholesky_inverse` is not supported with openvino backend"
     )
 
 

From bc3b37dd761c29f227c3022081942ab76a3e2f62 Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Sat, 9 Aug 2025 21:19:03 +0530
Subject: [PATCH 10/19] Modifying tests

---
 keras/src/ops/linalg.py      | 24 ++++++++++++++++--------
 keras/src/ops/linalg_test.py |  2 +-
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/keras/src/ops/linalg.py b/keras/src/ops/linalg.py
index 8dd9ff0b73f6..8b235da1b6bd 100644
--- a/keras/src/ops/linalg.py
+++ b/keras/src/ops/linalg.py
@@ -7,10 +7,14 @@
 
 
 class Cholesky(Operation):
-    def call(self, x, upper=False):
-        return _cholesky(x, upper)
+    def __init__(self, upper=False, name=None):
+        super().__init__(name=name)
+        self.upper = upper
+
+    def call(self, x):
+        return _cholesky(x, self.upper)
 
-    def compute_output_spec(self, x, upper=False):
+    def compute_output_spec(self, x):
         _assert_2d(x)
         _assert_square(x)
         return KerasTensor(x.shape, x.dtype)
@@ -29,7 +33,7 @@ def cholesky(x, upper=False):
         A tensor of shape `(..., M, M)` representing the Cholesky factor of `x`.
     """
     if any_symbolic_tensors((x,)):
-        return Cholesky().symbolic_call(x, upper=upper)
+        return Cholesky(upper=upper).symbolic_call(x)
     return _cholesky(x, upper=upper)
 
 
@@ -44,10 +48,14 @@ def _cholesky(x, upper=False):
 
 
 class CholeskyInverse(Operation):
-    def call(self, x, upper=False):
-        return _cholesky_inverse(x, upper)
+    def __init__(self, upper=False, name=None):
+        super().__init__(name=name)
+        self.upper = upper
+
+    def call(self, x):
+        return _cholesky_inverse(x, self.upper)
 
-    def compute_output_spec(self, x, upper=False):
+    def compute_output_spec(self, x):
         _assert_2d(x)
         _assert_square(x)
         return KerasTensor(x.shape, x.dtype)
@@ -69,7 +77,7 @@ def cholesky_inverse(x, upper=False):
         ValueError: If `x` is not a symmetric positive-definite matrix.
     """
     if any_symbolic_tensors((x,)):
-        return CholeskyInverse().symbolic_call(x, upper=upper)
+        return CholeskyInverse(upper=upper).symbolic_call(x)
     return _cholesky_inverse(x, upper=upper)
 
 
diff --git a/keras/src/ops/linalg_test.py b/keras/src/ops/linalg_test.py
index 2f7577161a27..0b3686989686 100644
--- a/keras/src/ops/linalg_test.py
+++ b/keras/src/ops/linalg_test.py
@@ -357,7 +357,7 @@ def test_cholesky(self):
         with self.assertRaises(ValueError):
             linalg.cholesky(x_non_psd)
         x = np.random.rand(4, 3, 3).astype("float32")
-        x_psd = np.matmul(x, x.transpose((0, 2, 1))) + 1e-5 * np.eye(3)
+        x_psd = np.matmul(x, x.transpose((0, 2, 1))) + 1e-5 * np.eye(3, dtype="float32")
 
         l_out = linalg.cholesky(x_psd, upper=False)
         l_expected = np.linalg.cholesky(x_psd)

From ca35406d7ff2edb1b24e3918bf9086df13933e0a Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Sun, 10 Aug 2025 22:55:36 +0530
Subject: [PATCH 11/19] Addressing the comments

---
 keras/src/backend/jax/linalg.py                        | 4 +---
 keras/src/backend/numpy/linalg.py                      | 4 +---
 keras/src/backend/openvino/excluded_concrete_tests.txt | 2 ++
 keras/src/backend/tensorflow/linalg.py                 | 2 +-
 keras/src/backend/torch/linalg.py                      | 8 ++------
 keras/src/ops/linalg.py                                | 4 ++--
 6 files changed, 9 insertions(+), 15 deletions(-)

diff --git a/keras/src/backend/jax/linalg.py b/keras/src/backend/jax/linalg.py
index acd992d5af44..12cce2ec022b 100644
--- a/keras/src/backend/jax/linalg.py
+++ b/keras/src/backend/jax/linalg.py
@@ -10,7 +10,7 @@
 
 
 def cholesky(a, upper=False):
-    out = jnp.linalg.cholesky(a)
+    out = jnp.linalg.cholesky(a, lower=not upper)
     try:
         # In eager mode, raise for nan to
         # achieve behavior consistency with numpy
@@ -23,8 +23,6 @@ def cholesky(a, upper=False):
     except jax.errors.TracerBoolConversionError:
         # Cannot raise for nan in tracing mode
         pass
-    if upper:
-        return jnp.swapaxes(out, -2, -1)
     return out
 
 
diff --git a/keras/src/backend/numpy/linalg.py b/keras/src/backend/numpy/linalg.py
index ea4bc16dac2e..c9a087d38cf9 100644
--- a/keras/src/backend/numpy/linalg.py
+++ b/keras/src/backend/numpy/linalg.py
@@ -7,9 +7,7 @@
 
 
 def cholesky(a, upper=False):
-    if upper:
-        return np.linalg.cholesky(a, upper=True)
-    return np.linalg.cholesky(a)
+    return np.linalg.cholesky(a, upper=upper)
 
 
 def cholesky_inverse(a, upper=False):
diff --git a/keras/src/backend/openvino/excluded_concrete_tests.txt b/keras/src/backend/openvino/excluded_concrete_tests.txt
index 7968a7931b89..01f413d65aee 100644
--- a/keras/src/backend/openvino/excluded_concrete_tests.txt
+++ b/keras/src/backend/openvino/excluded_concrete_tests.txt
@@ -289,3 +289,5 @@ TestMathErrors::test_istft_invalid_window_shape_2D_inputs
 TestMathErrors::test_stft_invalid_input_type
 TestMathErrors::test_stft_invalid_window
 TestMathErrors::test_stft_invalid_window_shape
+LinalgOpsCorrectnessTest::test_cholesky
+LinalgOpsCorrectnessTest::test_cholesky_inverse
diff --git a/keras/src/backend/tensorflow/linalg.py b/keras/src/backend/tensorflow/linalg.py
index eb22f3c3f480..5ed3cc3d36e9 100644
--- a/keras/src/backend/tensorflow/linalg.py
+++ b/keras/src/backend/tensorflow/linalg.py
@@ -12,7 +12,7 @@ def cholesky(a, upper=True):
     # tf.linalg.cholesky simply returns NaNs for non-positive definite matrices
     out = tf.debugging.check_numerics(out, "Cholesky")
     if upper:
-        return tf.linalg.matrix_transpose(out)
+        return tf.linalg.adjoint(out)
     return out
 
 
diff --git a/keras/src/backend/torch/linalg.py b/keras/src/backend/torch/linalg.py
index 2ccbc1c3e4e3..bae9733e36a4 100644
--- a/keras/src/backend/torch/linalg.py
+++ b/keras/src/backend/torch/linalg.py
@@ -8,15 +8,11 @@
 
 
 def cholesky(x, upper=False):
-    if upper:
-        return torch.linalg.cholesky(x, upper=True)
-    return torch.linalg.cholesky(x)
+    return torch.linalg.cholesky(x, upper=upper)
 
 
 def cholesky_inverse(x, upper=False):
-    if upper:
-        return torch.cholesky_inverse(x, upper=True)
-    return torch.cholesky_inverse(x)
+    return torch.cholesky_inverse(x, upper=upper)
 
 
 def det(x):
diff --git a/keras/src/ops/linalg.py b/keras/src/ops/linalg.py
index 8b235da1b6bd..05a49016871b 100644
--- a/keras/src/ops/linalg.py
+++ b/keras/src/ops/linalg.py
@@ -7,7 +7,7 @@
 
 
 class Cholesky(Operation):
-    def __init__(self, upper=False, name=None):
+    def __init__(self, upper=False, *, name=None):
         super().__init__(name=name)
         self.upper = upper
 
@@ -48,7 +48,7 @@ def _cholesky(x, upper=False):
 
 
 class CholeskyInverse(Operation):
-    def __init__(self, upper=False, name=None):
+    def __init__(self, upper=False, *, name=None):
         super().__init__(name=name)
         self.upper = upper
 

From 60ac5aa75b14cc37d0faea00d132e5ee82fa2f23 Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Mon, 11 Aug 2025 10:45:10 +0530
Subject: [PATCH 12/19] Modified test case for better tests

---
 keras/src/backend/jax/linalg.py |  2 +-
 keras/src/ops/linalg_test.py    | 22 +++++++++-------------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/keras/src/backend/jax/linalg.py b/keras/src/backend/jax/linalg.py
index 12cce2ec022b..83b89d5c5b25 100644
--- a/keras/src/backend/jax/linalg.py
+++ b/keras/src/backend/jax/linalg.py
@@ -10,7 +10,7 @@
 
 
 def cholesky(a, upper=False):
-    out = jnp.linalg.cholesky(a, lower=not upper)
+    out = jnp.linalg.cholesky(a, upper=upper)
     try:
         # In eager mode, raise for nan to
         # achieve behavior consistency with numpy
diff --git a/keras/src/ops/linalg_test.py b/keras/src/ops/linalg_test.py
index 0b3686989686..6c691588e3c4 100644
--- a/keras/src/ops/linalg_test.py
+++ b/keras/src/ops/linalg_test.py
@@ -370,22 +370,18 @@ def test_cholesky(self):
 
     def test_cholesky_inverse(self):
         x_np = np.random.rand(3, 3).astype("float32")
-        x_psd_np = np.matmul(x_np, x_np.T) + 1e-4 * np.eye(3, dtype="float32")
+        A = np.matmul(x_np, x_np.T) + 1e-5 * np.eye(3, dtype="float32")
         identity = np.eye(3, dtype="float32")
 
-        l_factor_np = np.linalg.cholesky(x_psd_np)
-        x_inv_from_l = linalg.cholesky_inverse(l_factor_np, upper=False)
-        reconstructed_from_l = ops.matmul(x_psd_np, x_inv_from_l)
-        self.assertAllClose(
-            reconstructed_from_l, identity, atol=1e-4, rtol=1e-4
-        )
+        L = np.linalg.cholesky(A)
+        A_inv_from_l = ops.linalg.cholesky_inverse(L, upper=False)
+        reconstructed_from_l = ops.matmul(A, A_inv_from_l)
+        self.assertAllClose(reconstructed_from_l, identity, atol=1e-4)
 
-        u_factor_np = l_factor_np.T
-        x_inv_from_u = linalg.cholesky_inverse(u_factor_np, upper=True)
-        reconstructed_from_u = ops.matmul(x_psd_np, x_inv_from_u)
-        self.assertAllClose(
-            reconstructed_from_u, identity, atol=1e-4, rtol=1e-4
-        )
+        U = np.linalg.cholesky(A, upper=True)
+        A_inv_from_u = ops.linalg.cholesky_inverse(U, upper=True)
+        reconstructed_from_u = ops.matmul(A, A_inv_from_u)
+        self.assertAllClose(reconstructed_from_u, identity, atol=1e-4)
 
     def test_det(self):
         x = np.random.rand(4, 3, 3)

From 3edaeb7781c8b5773d3a7425ea6a26ed0dc6ce2b Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Mon, 11 Aug 2025 13:07:04 +0530
Subject: [PATCH 13/19] Changed the test case

---
 keras/api/_tf_keras/keras/ops/__init__.py     |  1 +
 .../_tf_keras/keras/ops/linalg/__init__.py    |  1 +
 keras/api/ops/__init__.py                     |  1 +
 keras/api/ops/linalg/__init__.py              |  1 +
 keras/src/backend/openvino/linalg.py          |  4 +-
 keras/src/backend/tensorflow/linalg.py        |  4 +-
 keras/src/ops/linalg.py                       |  4 +-
 keras/src/ops/linalg_test.py                  | 42 ++++++++++++-------
 8 files changed, 38 insertions(+), 20 deletions(-)

diff --git a/keras/api/_tf_keras/keras/ops/__init__.py b/keras/api/_tf_keras/keras/ops/__init__.py
index 0d4037717cc2..5ff4561e9d5e 100644
--- a/keras/api/_tf_keras/keras/ops/__init__.py
+++ b/keras/api/_tf_keras/keras/ops/__init__.py
@@ -32,6 +32,7 @@
 from keras.src.ops.core import while_loop as while_loop
 from keras.src.ops.einops import rearrange as rearrange
 from keras.src.ops.linalg import cholesky as cholesky
+from keras.src.ops.linalg import cholesky_inverse as cholesky_inverse
 from keras.src.ops.linalg import det as det
 from keras.src.ops.linalg import eig as eig
 from keras.src.ops.linalg import eigh as eigh
diff --git a/keras/api/_tf_keras/keras/ops/linalg/__init__.py b/keras/api/_tf_keras/keras/ops/linalg/__init__.py
index bc091ea766a4..0c96c3bbb8dc 100644
--- a/keras/api/_tf_keras/keras/ops/linalg/__init__.py
+++ b/keras/api/_tf_keras/keras/ops/linalg/__init__.py
@@ -5,6 +5,7 @@
 """
 
 from keras.src.ops.linalg import cholesky as cholesky
+from keras.src.ops.linalg import cholesky_inverse as cholesky_inverse
 from keras.src.ops.linalg import det as det
 from keras.src.ops.linalg import eig as eig
 from keras.src.ops.linalg import eigh as eigh
diff --git a/keras/api/ops/__init__.py b/keras/api/ops/__init__.py
index 0d4037717cc2..5ff4561e9d5e 100644
--- a/keras/api/ops/__init__.py
+++ b/keras/api/ops/__init__.py
@@ -32,6 +32,7 @@
 from keras.src.ops.core import while_loop as while_loop
 from keras.src.ops.einops import rearrange as rearrange
 from keras.src.ops.linalg import cholesky as cholesky
+from keras.src.ops.linalg import cholesky_inverse as cholesky_inverse
 from keras.src.ops.linalg import det as det
 from keras.src.ops.linalg import eig as eig
 from keras.src.ops.linalg import eigh as eigh
diff --git a/keras/api/ops/linalg/__init__.py b/keras/api/ops/linalg/__init__.py
index bc091ea766a4..0c96c3bbb8dc 100644
--- a/keras/api/ops/linalg/__init__.py
+++ b/keras/api/ops/linalg/__init__.py
@@ -5,6 +5,7 @@
 """
 
 from keras.src.ops.linalg import cholesky as cholesky
+from keras.src.ops.linalg import cholesky_inverse as cholesky_inverse
 from keras.src.ops.linalg import det as det
 from keras.src.ops.linalg import eig as eig
 from keras.src.ops.linalg import eigh as eigh
diff --git a/keras/src/backend/openvino/linalg.py b/keras/src/backend/openvino/linalg.py
index 31ae0abe37c6..948c4d480144 100644
--- a/keras/src/backend/openvino/linalg.py
+++ b/keras/src/backend/openvino/linalg.py
@@ -1,12 +1,12 @@
 def cholesky(a, upper=False):
     raise NotImplementedError(
-        "`cholesky` is not supported with openvino backend"
+        "`cholesky` is not supported with openvino backend."
     )
 
 
 def cholesky_inverse(a, upper=False):
     raise NotImplementedError(
-        "`cholesky_inverse` is not supported with openvino backend"
+        "`cholesky_inverse` is not supported with openvino backend."
     )
 
 
diff --git a/keras/src/backend/tensorflow/linalg.py b/keras/src/backend/tensorflow/linalg.py
index 5ed3cc3d36e9..f8fdc6501c72 100644
--- a/keras/src/backend/tensorflow/linalg.py
+++ b/keras/src/backend/tensorflow/linalg.py
@@ -7,7 +7,7 @@
 from keras.src.backend.tensorflow.core import convert_to_tensor
 
 
-def cholesky(a, upper=True):
+def cholesky(a, upper=False):
     out = tf.linalg.cholesky(a)
     # tf.linalg.cholesky simply returns NaNs for non-positive definite matrices
     out = tf.debugging.check_numerics(out, "Cholesky")
@@ -16,7 +16,7 @@ def cholesky(a, upper=True):
     return out
 
 
-def cholesky_inverse(a, upper=True):
+def cholesky_inverse(a, upper=False):
     identity = tf.eye(num_rows=tf.shape(a)[-1], dtype=a.dtype)
     if upper:
         u_inv = tf.linalg.triangular_solve(a, identity, lower=False)
diff --git a/keras/src/ops/linalg.py b/keras/src/ops/linalg.py
index 05a49016871b..c294454dd5b2 100644
--- a/keras/src/ops/linalg.py
+++ b/keras/src/ops/linalg.py
@@ -61,7 +61,9 @@ def compute_output_spec(self, x):
         return KerasTensor(x.shape, x.dtype)
 
 
-@keras_export(["keras.ops.cholesky_inverse", "keras.ops.linalg.cholesky_inverse"])
+@keras_export(
+    ["keras.ops.cholesky_inverse", "keras.ops.linalg.cholesky_inverse"]
+)
 def cholesky_inverse(x, upper=False):
     """Computes the inverse of a symmetric positive-definite matrix.
 
diff --git a/keras/src/ops/linalg_test.py b/keras/src/ops/linalg_test.py
index 6c691588e3c4..a7fcad6c507a 100644
--- a/keras/src/ops/linalg_test.py
+++ b/keras/src/ops/linalg_test.py
@@ -5,6 +5,7 @@
 from keras.src import ops
 from keras.src import testing
 from keras.src.backend.common.keras_tensor import KerasTensor
+from keras.src.backend.numpy import linalg as numpy_linalg_backend
 from keras.src.ops import linalg
 from keras.src.testing.test_utils import named_product
 
@@ -357,7 +358,9 @@ def test_cholesky(self):
         with self.assertRaises(ValueError):
             linalg.cholesky(x_non_psd)
         x = np.random.rand(4, 3, 3).astype("float32")
-        x_psd = np.matmul(x, x.transpose((0, 2, 1))) + 1e-5 * np.eye(3, dtype="float32")
+        x_psd = np.matmul(x, x.transpose((0, 2, 1))) + 1e-5 * np.eye(
+            3, dtype="float32"
+        )
 
         l_out = linalg.cholesky(x_psd, upper=False)
         l_expected = np.linalg.cholesky(x_psd)
@@ -367,21 +370,30 @@ def test_cholesky(self):
         u_expected = l_expected.transpose((0, 2, 1))
         self.assertAllClose(u_out, u_expected, atol=1e-4)
 
-
     def test_cholesky_inverse(self):
-        x_np = np.random.rand(3, 3).astype("float32")
-        A = np.matmul(x_np, x_np.T) + 1e-5 * np.eye(3, dtype="float32")
-        identity = np.eye(3, dtype="float32")
-
-        L = np.linalg.cholesky(A)
-        A_inv_from_l = ops.linalg.cholesky_inverse(L, upper=False)
-        reconstructed_from_l = ops.matmul(A, A_inv_from_l)
-        self.assertAllClose(reconstructed_from_l, identity, atol=1e-4)
-
-        U = np.linalg.cholesky(A, upper=True)
-        A_inv_from_u = ops.linalg.cholesky_inverse(U, upper=True)
-        reconstructed_from_u = ops.matmul(A, A_inv_from_u)
-        self.assertAllClose(reconstructed_from_u, identity, atol=1e-4)
+        x = np.random.rand(4, 3, 3).astype("float32")
+        x_psd_batch = np.matmul(x, x.transpose((0, 2, 1))) + 1e-5 * np.eye(
+            3, dtype="float32"
+        )
+
+        for i in range(x_psd_batch.shape[0]):
+            x_psd = x_psd_batch[i]
+
+            l_out = linalg.cholesky_inverse(x_psd, upper=False)
+            l_expected = numpy_linalg_backend.cholesky_inverse(
+                x_psd, upper=False
+            )
+            self.assertAllClose(
+                l_out, l_expected, atol=1e-4, msg=f"Matrix {i} (lower) failed"
+            )
+
+            u_out = linalg.cholesky_inverse(x_psd, upper=True)
+            u_expected = numpy_linalg_backend.cholesky_inverse(
+                x_psd, upper=True
+            )
+            self.assertAllClose(
+                u_out, u_expected, atol=1e-4, msg=f"Matrix {i} (upper) failed"
+            )
 
     def test_det(self):
         x = np.random.rand(4, 3, 3)

From d190274b0b6b0302f8584d228925cc5d9c67aaf1 Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Mon, 11 Aug 2025 13:21:52 +0530
Subject: [PATCH 14/19] modified the test case

---
 keras/src/ops/linalg_test.py | 40 +++++++++++++++---------------------
 1 file changed, 17 insertions(+), 23 deletions(-)

diff --git a/keras/src/ops/linalg_test.py b/keras/src/ops/linalg_test.py
index a7fcad6c507a..6c2b7cb12410 100644
--- a/keras/src/ops/linalg_test.py
+++ b/keras/src/ops/linalg_test.py
@@ -5,7 +5,6 @@
 from keras.src import ops
 from keras.src import testing
 from keras.src.backend.common.keras_tensor import KerasTensor
-from keras.src.backend.numpy import linalg as numpy_linalg_backend
 from keras.src.ops import linalg
 from keras.src.testing.test_utils import named_product
 
@@ -371,29 +370,24 @@ def test_cholesky(self):
         self.assertAllClose(u_out, u_expected, atol=1e-4)
 
     def test_cholesky_inverse(self):
-        x = np.random.rand(4, 3, 3).astype("float32")
-        x_psd_batch = np.matmul(x, x.transpose((0, 2, 1))) + 1e-5 * np.eye(
-            3, dtype="float32"
+        input = np.array(
+            [
+                [4.0, 12.0, -16.0],
+                [12.0, 37.0, -43.0],
+                [-16.0, -43.0, 98.0],
+            ],
+            dtype="float32",
         )
-
-        for i in range(x_psd_batch.shape[0]):
-            x_psd = x_psd_batch[i]
-
-            l_out = linalg.cholesky_inverse(x_psd, upper=False)
-            l_expected = numpy_linalg_backend.cholesky_inverse(
-                x_psd, upper=False
-            )
-            self.assertAllClose(
-                l_out, l_expected, atol=1e-4, msg=f"Matrix {i} (lower) failed"
-            )
-
-            u_out = linalg.cholesky_inverse(x_psd, upper=True)
-            u_expected = numpy_linalg_backend.cholesky_inverse(
-                x_psd, upper=True
-            )
-            self.assertAllClose(
-                u_out, u_expected, atol=1e-4, msg=f"Matrix {i} (upper) failed"
-            )
+        expected_output = np.array(
+            [
+                [6.910159e-02, -2.129242e-03, 5.346869e-05],
+                [-2.129242e-03, 8.710913e-04, 1.210081e-04],
+                [5.346869e-05, 1.210081e-04, 1.041233e-04],
+            ],
+            dtype="float32",
+        )
+        op_output = linalg.cholesky_inverse(input)
+        self.assertAllClose(op_output, expected_output, atol=1e-5)
 
     def test_det(self):
         x = np.random.rand(4, 3, 3)

From b1f79d1f5cc025d9a3546807b05ee98a5ddf1338 Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Mon, 11 Aug 2025 13:53:40 +0530
Subject: [PATCH 15/19] Added more tests

---
 keras/src/ops/linalg_test.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/keras/src/ops/linalg_test.py b/keras/src/ops/linalg_test.py
index 6c2b7cb12410..1cac3b7b95b7 100644
--- a/keras/src/ops/linalg_test.py
+++ b/keras/src/ops/linalg_test.py
@@ -369,8 +369,13 @@ def test_cholesky(self):
         u_expected = l_expected.transpose((0, 2, 1))
         self.assertAllClose(u_out, u_expected, atol=1e-4)
 
-    def test_cholesky_inverse(self):
-        input = np.array(
+    @parameterized.named_parameters(
+        ("lower", False),
+        ("upper", True),
+    )
+    def test_cholesky_inverse(self, upper):
+        """Tests cholesky_inverse for both lower and upper computations."""
+        input_matrix = np.array(
             [
                 [4.0, 12.0, -16.0],
                 [12.0, 37.0, -43.0],
@@ -386,7 +391,7 @@ def test_cholesky_inverse(self):
             ],
             dtype="float32",
         )
-        op_output = linalg.cholesky_inverse(input)
+        op_output = linalg.cholesky_inverse(input_matrix, upper=upper)
         self.assertAllClose(op_output, expected_output, atol=1e-5)
 
     def test_det(self):

From d2591824a210dbfb06c6ad8d27eed21794c66210 Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Mon, 11 Aug 2025 14:09:43 +0530
Subject: [PATCH 16/19] Added more tests

---
 keras/src/ops/linalg_test.py | 25 +++++++++++--------------
 1 file changed, 11 insertions(+), 14 deletions(-)

diff --git a/keras/src/ops/linalg_test.py b/keras/src/ops/linalg_test.py
index 1cac3b7b95b7..e9a6bd30217a 100644
--- a/keras/src/ops/linalg_test.py
+++ b/keras/src/ops/linalg_test.py
@@ -370,12 +370,11 @@ def test_cholesky(self):
         self.assertAllClose(u_out, u_expected, atol=1e-4)
 
     @parameterized.named_parameters(
-        ("lower", False),
-        ("upper", True),
+        {"testcase_name": "lower", "upper": False},
+        {"testcase_name": "upper", "upper": True},
     )
     def test_cholesky_inverse(self, upper):
-        """Tests cholesky_inverse for both lower and upper computations."""
-        input_matrix = np.array(
+        A = np.array(
             [
                 [4.0, 12.0, -16.0],
                 [12.0, 37.0, -43.0],
@@ -383,16 +382,14 @@ def test_cholesky_inverse(self, upper):
             ],
             dtype="float32",
         )
-        expected_output = np.array(
-            [
-                [6.910159e-02, -2.129242e-03, 5.346869e-05],
-                [-2.129242e-03, 8.710913e-04, 1.210081e-04],
-                [5.346869e-05, 1.210081e-04, 1.041233e-04],
-            ],
-            dtype="float32",
-        )
-        op_output = linalg.cholesky_inverse(input_matrix, upper=upper)
-        self.assertAllClose(op_output, expected_output, atol=1e-5)
+        if upper:
+            factor = np.linalg.cholesky(A, upper=True)
+        else:
+            factor = np.linalg.cholesky(A)
+
+        expected_inverse = np.linalg.inv(A)
+        output_inverse = linalg.cholesky_inverse(factor, upper=upper)
+        self.assertAllClose(output_inverse, expected_inverse, atol=1e-5)
 
     def test_det(self):
         x = np.random.rand(4, 3, 3)

From 180d5dbb0f8e0de388aa921a0fb44dbcad104cbc Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Mon, 11 Aug 2025 14:15:39 +0530
Subject: [PATCH 17/19] correcting the test case

---
 keras/src/ops/linalg_test.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/keras/src/ops/linalg_test.py b/keras/src/ops/linalg_test.py
index e9a6bd30217a..0c0137ce687f 100644
--- a/keras/src/ops/linalg_test.py
+++ b/keras/src/ops/linalg_test.py
@@ -387,7 +387,15 @@ def test_cholesky_inverse(self, upper):
         else:
             factor = np.linalg.cholesky(A)
 
-        expected_inverse = np.linalg.inv(A)
+        expected_inverse = np.array(
+            [
+                [49.36111, -13.555555, 2.111111],
+                [-13.555555, 3.777778, -0.555556],
+                [2.111111, -0.555556, 0.111111],
+            ],
+            dtype="float32",
+        )
+
         output_inverse = linalg.cholesky_inverse(factor, upper=upper)
         self.assertAllClose(output_inverse, expected_inverse, atol=1e-5)
 

From 8fd86450f03c11c0a04acecbad652a43073ca731 Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Tue, 12 Aug 2025 09:43:29 +0530
Subject: [PATCH 18/19] Modifying cholesky_inverse functions

---
 keras/src/backend/jax/linalg.py        | 7 +++----
 keras/src/backend/numpy/linalg.py      | 7 +++----
 keras/src/backend/tensorflow/linalg.py | 7 +++----
 3 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/keras/src/backend/jax/linalg.py b/keras/src/backend/jax/linalg.py
index 83b89d5c5b25..f36b80dcb09d 100644
--- a/keras/src/backend/jax/linalg.py
+++ b/keras/src/backend/jax/linalg.py
@@ -28,12 +28,11 @@ def cholesky(a, upper=False):
 
 def cholesky_inverse(a, upper=False):
     identity = jnp.eye(a.shape[-1], dtype=a.dtype)
+    inv_chol = solve_triangular(a, identity, lower=not upper)
     if upper:
-        u_inv = solve_triangular(a, identity, lower=False)
-        a_inv = jnp.matmul(u_inv, jnp.transpose(u_inv))
+        a_inv = jnp.matmul(inv_chol, jnp.transpose(inv_chol))
     else:
-        l_inv = solve_triangular(a, identity, lower=True)
-        a_inv = jnp.matmul(jnp.transpose(l_inv), l_inv)
+        a_inv = jnp.matmul(jnp.transpose(inv_chol), inv_chol)
     return a_inv
 
 
diff --git a/keras/src/backend/numpy/linalg.py b/keras/src/backend/numpy/linalg.py
index c9a087d38cf9..9fe27f6aac11 100644
--- a/keras/src/backend/numpy/linalg.py
+++ b/keras/src/backend/numpy/linalg.py
@@ -12,12 +12,11 @@ def cholesky(a, upper=False):
 
 def cholesky_inverse(a, upper=False):
     identity = np.eye(a.shape[-1], dtype=a.dtype)
+    inv_chol = solve_triangular(a, identity, lower=not upper)
     if upper:
-        u_inv = solve_triangular(a, identity, lower=False)
-        a_inv = np.matmul(u_inv, u_inv.T)
+        a_inv = np.matmul(inv_chol, inv_chol.T)
     else:
-        l_inv = solve_triangular(a, identity, lower=True)
-        a_inv = np.matmul(l_inv.T, l_inv)
+        a_inv = np.matmul(inv_chol.T, inv_chol)
     return a_inv
 
 
diff --git a/keras/src/backend/tensorflow/linalg.py b/keras/src/backend/tensorflow/linalg.py
index f8fdc6501c72..9a1f1b615249 100644
--- a/keras/src/backend/tensorflow/linalg.py
+++ b/keras/src/backend/tensorflow/linalg.py
@@ -18,12 +18,11 @@ def cholesky(a, upper=False):
 
 def cholesky_inverse(a, upper=False):
     identity = tf.eye(num_rows=tf.shape(a)[-1], dtype=a.dtype)
+    inv_chol = tf.linalg.triangular_solve(a, identity, lower=not upper)
     if upper:
-        u_inv = tf.linalg.triangular_solve(a, identity, lower=False)
-        a_inv = tf.matmul(u_inv, u_inv, transpose_b=True)
+        a_inv = tf.matmul(inv_chol, inv_chol, transpose_b=True)
     else:
-        l_inv = tf.linalg.triangular_solve(a, identity, lower=True)
-        a_inv = tf.matmul(l_inv, l_inv, transpose_a=True)
+        a_inv = tf.matmul(inv_chol, inv_chol, transpose_a=True)
     return a_inv
 
 

From a70b69ccfb6fdc05d22a3a33e502eb845537ec14 Mon Sep 17 00:00:00 2001
From: suhana <suhanaaa@google.com>
Date: Tue, 12 Aug 2025 13:00:53 +0530
Subject: [PATCH 19/19] Fixed nit comments

---
 keras/src/ops/linalg_test.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/keras/src/ops/linalg_test.py b/keras/src/ops/linalg_test.py
index 0c0137ce687f..5ae00af915c8 100644
--- a/keras/src/ops/linalg_test.py
+++ b/keras/src/ops/linalg_test.py
@@ -356,6 +356,7 @@ def test_cholesky(self):
         x_non_psd = np.random.rand(4, 3, 3).astype("float32")
         with self.assertRaises(ValueError):
             linalg.cholesky(x_non_psd)
+
         x = np.random.rand(4, 3, 3).astype("float32")
         x_psd = np.matmul(x, x.transpose((0, 2, 1))) + 1e-5 * np.eye(
             3, dtype="float32"