Raise if Cholesky output has nans when on_error="raise"

Author: jessegrabowski

pytensor/link/numba/dispatch/slinalg.py

@@ -292,7 +292,7 @@ def solve_triangular(a, b):
         res = _solve_triangular(a, b, trans, lower, unit_diagonal)
         if check_finite:
             if np.any(np.bitwise_or(np.isinf(res), np.isnan(res))):
-                raise ValueError(
+                raise np.linalg.LinAlgError(
                     "Non-numeric values (nan or inf) returned by solve_triangular"
                 )
         return res
@@ -349,11 +349,17 @@ def impl(A, lower=0, overwrite_a=False, check_finite=True):
 def numba_funcify_Cholesky(op, node, **kwargs):
     lower = op.lower
     overwrite_a = False
-    on_error = op.on_error
+    check_finite = op.on_error == "raise"
 
     @numba_basic.numba_njit(inline="always")
     def nb_cholesky(a):
-        res = _cholesky(a, lower, overwrite_a, on_error)
+        res = _cholesky(a, lower, overwrite_a, check_finite)
+        if check_finite:
+            if np.any(np.bitwise_or(np.isinf(res), np.isnan(res))):
+                raise np.linalg.LinAlgError(
+                    "Non-numeric values (nan or inf) returned by cholesky"
+                )
+
         return res
 
     return nb_cholesky

tests/link/numba/test_slinalg.py

@@ -102,7 +102,8 @@ def test_solve_triangular_raises_on_nan_inf(value):
     b = np.full((5, 1), value)
 
     with pytest.raises(
-        ValueError, match=re.escape("Non-numeric values (nan or inf) returned ")
+        np.linalg.LinAlgError,
+        match=re.escape("Non-numeric values (nan or inf) returned "),
     ):
         f(A_tri, b)
 
@@ -136,7 +137,7 @@ def test_numba_Cholesky_raises_on_nan():
     g = pt.linalg.cholesky(x, on_error="raise")
     f = pytensor.function([x], g, mode="NUMBA")
 
-    with pytest.raises(ValueError, match=r"Non-numeric values"):
+    with pytest.raises(np.linalg.LinAlgError, match=r"Non-numeric values"):
         f(test_value)