Add numba dispatch for lu_factor

Author: jessegrabowski

pytensor/link/numba/dispatch/slinalg.py

@@ -478,7 +478,8 @@ def impl(A: np.ndarray, A_norm: float, norm: str) -> tuple[np.ndarray, int]:
 
 def _getrf(A, overwrite_a=False) -> tuple[np.ndarray, np.ndarray, int]:
     """
-    Placeholder for LU factorization; used by linalg.solve.
+    Underlying LAPACK function used for LU factorization. Compared to scipy.linalg.lu_factorize, this function also
+    returns an info code with diagnostic information.
     """
     getrf = scipy.linalg.get_lapack_funcs("getrf", (A,))
     A_copy, ipiv, info = getrf(A, overwrite_a=overwrite_a)
@@ -519,6 +520,29 @@ def impl(
     return impl
 
 
+def _lu_factor(A, overwrite_a=False) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Thin wrapper around scipy.linalg.lu_factor. Used as an overload target to avoid side-effects on users who import
+    Pytensor.
+    """
+    return linalg.lu_factor(A, overwrite_a=overwrite_a)
+
+
+@overload(_lu_factor)
+def lu_factor_impl(
+    A: np.ndarray, overwrite_a: bool = False
+) -> Callable[[np.ndarray, bool], tuple[np.ndarray, np.ndarray]]:
+    ensure_lapack()
+    _check_scipy_linalg_matrix(A, "lu_factor")
+
+    def impl(A: np.ndarray, overwrite_a: bool = False) -> tuple[np.ndarray, np.ndarray]:
+        A_copy, IPIV, INFO = _getrf(A, overwrite_a=overwrite_a)
+        _solve_check(int_ptr_to_val(INFO), 0)
+        return A_copy, IPIV
+
+    return impl
+
+
 def _lu_1(
     a: np.ndarray,
     permute_l: bool,

tests/link/numba/test_slinalg.py

@@ -521,10 +521,31 @@ def test_numba_lu(permute_l, p_indices, shape: tuple[int]):
 
     else:
         # compare_numba_and_py fails: NotImplementedError: Non-jitted BlockwiseWithCoreShape not implemented
-        nb_out = f(A_val.copy())
+        pt_out = f(A_val.copy())
         sp_out = scipy_linalg.lu(
             A_val.copy(), permute_l=permute_l, p_indices=p_indices, check_finite=False
         )
 
-        for a, b in zip(nb_out, sp_out, strict=True):
+        for a, b in zip(pt_out, sp_out, strict=True):
+            np.testing.assert_allclose(a, b)
+
+
+@pytest.mark.parametrize("shape", [(3, 5, 5), (5, 5)], ids=["batched", "not_batched"])
+def test_numba_lu_factor(shape: tuple[int]):
+    rng = np.random.default_rng(utt.fetch_seed())
+    A = pt.tensor("A", shape=shape, dtype=config.floatX)
+    out = pt.linalg.lu_factor(A)
+
+    A_val = rng.normal(size=shape).astype(config.floatX)
+    f = pytensor.function([A], out, mode="NUMBA")
+
+    if len(shape) == 2:
+        compare_numba_and_py([A], out, [A_val], inplace=True)
+    else:
+        pt_out = f(A_val.copy())
+        sp_out = np.vectorize(scipy_linalg.lu_factor, signature="(n,n)->(n,n),(n)")(
+            A_val.copy()
+        )
+
+        for a, b in zip(pt_out, sp_out, strict=True):
             np.testing.assert_allclose(a, b)