Validate inputs and outputs of numba cholesky function

Author: jessegrabowski

pytensor/link/numba/dispatch/slinalg.py

@@ -340,7 +340,7 @@ def impl(A, lower=0, overwrite_a=False, check_finite=True):
             INFO,
         )
 
-        return A_copy
+        return A_copy, int_ptr_to_val(INFO)
 
     return impl
 
@@ -349,16 +349,30 @@ def impl(A, lower=0, overwrite_a=False, check_finite=True):
 def numba_funcify_Cholesky(op, node, **kwargs):
     lower = op.lower
     overwrite_a = False
-    check_finite = op.on_error == "raise"
+    check_finite = op.check_finite
+    on_error = op.on_error
 
     @numba_basic.numba_njit(inline="always")
     def nb_cholesky(a):
-        res = _cholesky(a, lower, overwrite_a, check_finite)
         if check_finite:
-            if np.any(np.bitwise_or(np.isinf(res), np.isnan(res))):
+            if np.any(np.bitwise_or(np.isinf(a), np.isnan(a))):
                 raise np.linalg.LinAlgError(
-                    "Non-numeric values (nan or inf) returned by cholesky"
+                    "Non-numeric values (nan or inf) found in input to cholesky"
                 )
+        res, info = _cholesky(a, lower, overwrite_a, check_finite)
+
+        if on_error == "raise":
+            if info > 0:
+                raise np.linalg.LinAlgError(
+                    "Input to cholesky is not positive definite"
+                )
+            if info < 0:
+                raise ValueError(
+                    'LAPACK reported an illegal value in input on entry to "POTRF."'
+                )
+        else:
+            if info != 0:
+                res = np.full_like(res, np.nan)
 
         return res
 

pytensor/tensor/slinalg.py

@@ -51,9 +51,10 @@ class Cholesky(Op):
     __props__ = ("lower", "destructive", "on_error")
     gufunc_signature = "(m,m)->(m,m)"
 
-    def __init__(self, *, lower=True, on_error="raise"):
+    def __init__(self, *, lower=True, check_finite=True, on_error="raise"):
         self.lower = lower
         self.destructive = False
+        self.check_finite = check_finite
         if on_error not in ("raise", "nan"):
             raise ValueError('on_error must be one of "raise" or ""nan"')
         self.on_error = on_error
@@ -70,7 +71,9 @@ def perform(self, node, inputs, outputs):
         x = inputs[0]
         z = outputs[0]
         try:
-            z[0] = scipy.linalg.cholesky(x, lower=self.lower).astype(x.dtype)
+            z[0] = scipy.linalg.cholesky(
+                x, lower=self.lower, check_finite=self.check_finite
+            ).astype(x.dtype)
         except scipy.linalg.LinAlgError:
             if self.on_error == "raise":
                 raise
@@ -129,8 +132,10 @@ def conjugate_solve_triangular(outer, inner):
             return [grad]
 
 
-def cholesky(x, lower=True, on_error="raise"):
-    return Blockwise(Cholesky(lower=lower, on_error=on_error))(x)
+def cholesky(x, lower=True, on_error="raise", check_finite=True):
+    return Blockwise(
+        Cholesky(lower=lower, on_error=on_error, check_finite=check_finite)
+    )(x)
 
 
 class SolveBase(Op):

tests/link/numba/test_slinalg.py

@@ -128,19 +128,36 @@ def test_numba_Cholesky(lower):
     )
 
 
-def test_numba_Cholesky_raises_on_nan():
+def test_numba_Cholesky_raises_on_nan_input():
     test_value = rng.random(size=(3, 3)).astype(config.floatX)
     test_value[0, 0] = np.nan
 
     x = pt.tensor(dtype=config.floatX, shape=(3, 3))
     x = x.T.dot(x)
-    g = pt.linalg.cholesky(x, on_error="raise")
+    g = pt.linalg.cholesky(x, check_finite=True)
     f = pytensor.function([x], g, mode="NUMBA")
 
     with pytest.raises(np.linalg.LinAlgError, match=r"Non-numeric values"):
         f(test_value)
 
 
+@pytest.mark.parametrize("on_error", ["nan", "raise"])
+def test_numba_Cholesky_raise_on(on_error):
+    test_value = rng.random(size=(3, 3)).astype(config.floatX)
+
+    x = pt.tensor(dtype=config.floatX, shape=(3, 3))
+    g = pt.linalg.cholesky(x, on_error=on_error)
+    f = pytensor.function([x], g, mode="NUMBA")
+
+    if on_error == "raise":
+        with pytest.raises(
+            np.linalg.LinAlgError, match=r"Input to cholesky is not positive definite"
+        ):
+            f(test_value)
+    else:
+        assert np.all(np.isnan(f(test_value)))
+
+
 def test_block_diag():
     A = pt.matrix("A")
     B = pt.matrix("B")