Use lapack func instead of scipy.linalg.cholesky

pytensor/tensor/slinalg.py

@@ -7,6 +7,7 @@
 import numpy as np
 import scipy.linalg as scipy_linalg
 from numpy.exceptions import ComplexWarning
+from scipy.linalg._misc import _datacopied
 
 import pytensor
 import pytensor.tensor as pt
@@ -37,7 +38,7 @@ def __init__(
         self,
         *,
         lower: bool = True,
-        check_finite: bool = True,
+        check_finite: bool = False,
         on_error: Literal["raise", "nan"] = "raise",
         overwrite_a: bool = False,
     ):
@@ -64,21 +65,52 @@ def make_node(self, x):
         dtype = scipy_linalg.cholesky(np.eye(1, dtype=x.type.dtype)).dtype
         return Apply(self, [x], [tensor(shape=x.type.shape, dtype=dtype)])
 
+    def _cholesky(
+        self, a, lower=False, overwrite_a=False, clean=True, check_finite=False
+    ):
+        a1 = np.asarray_chkfinite(a) if check_finite else np.asarray(a)
+
+        # Squareness check
+        if a1.shape[0] != a1.shape[1]:
+            raise ValueError(
+                "Input array is expected to be square but has "
+                f"the shape: {a1.shape}."
+            )
+
+        # Quick return for square empty array
+        if a1.size == 0:
+            dt = self._cholesky(np.eye(1, dtype=a1.dtype)).dtype
+            return np.empty_like(a1, dtype=dt), lower
+
+        overwrite_a = overwrite_a or _datacopied(a1, a)
+        (potrf,) = scipy_linalg.get_lapack_funcs(("potrf",), (a1,))
+        c, info = potrf(a1, lower=lower, overwrite_a=overwrite_a, clean=clean)
+        if info > 0:
+            raise scipy_linalg.LinAlgError(
+                f"{info}-th leading minor of the array is not positive definite"
+            )
+        if info < 0:
+            raise ValueError(
+                f"LAPACK reported an illegal value in {-info}-th argument "
+                f'on entry to "POTRF".'
+            )
+        return c
+
     def perform(self, node, inputs, outputs):
         [x] = inputs
         [out] = outputs
         try:
             # Scipy cholesky only makes use of overwrite_a when it is F_CONTIGUOUS
             # If we have a `C_CONTIGUOUS` array we transpose to benefit from it
             if self.overwrite_a and x.flags["C_CONTIGUOUS"]:
-                out[0] = scipy_linalg.cholesky(
+                out[0] = self._cholesky(
                     x.T,
                     lower=not self.lower,
                     check_finite=self.check_finite,
                     overwrite_a=True,
                 ).T
             else:
-                out[0] = scipy_linalg.cholesky(
+                out[0] = self._cholesky(
                     x,
                     lower=self.lower,
                     check_finite=self.check_finite,
@@ -201,7 +233,9 @@ def cholesky(
 
     """
 
-    return Blockwise(Cholesky(lower=lower, on_error=on_error))(x)
+    return Blockwise(
+        Cholesky(lower=lower, on_error=on_error, check_finite=check_finite)
+    )(x)
 
 
 class SolveBase(Op):

tests/tensor/test_slinalg.py

@@ -74,6 +74,16 @@ def test_cholesky():
     check_upper_triangular(pd, ch_f)
 
 
+def test_cholesky_performance(benchmark):
+    rng = np.random.default_rng(utt.fetch_seed())
+    r = rng.standard_normal((10, 10)).astype(config.floatX)
+    pd = np.dot(r, r.T)
+    x = matrix()
+    chol = cholesky(x)
+    ch_f = function([x], chol)
+    benchmark(ch_f, pd)
+
+
 def test_cholesky_indef():
     x = matrix()
     mat = np.array([[1, 0.2], [0.2, -2]]).astype(config.floatX)