Numba linalg: Handle empty inputs

Author: ricardoV94

pytensor/link/numba/dispatch/linalg/decomposition/lu_factor.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 from numba.core.extending import overload
+from numba.core.types import Float
 from numba.np.linalg import _copy_to_fortran_order, ensure_lapack
 from scipy import linalg
 
@@ -35,7 +36,7 @@ def getrf_impl(
     A: np.ndarray, overwrite_a: bool = False
 ) -> Callable[[np.ndarray, bool], tuple[np.ndarray, np.ndarray, int]]:
     ensure_lapack()
-    _check_linalg_matrix(A, ndim=2, func_name="getrf")
+    _check_linalg_matrix(A, ndim=2, dtype=Float, func_name="getrf")
     dtype = A.dtype
     numba_getrf = _LAPACK().numba_xgetrf(dtype)
 
@@ -75,7 +76,7 @@ def lu_factor_impl(
     A: np.ndarray, overwrite_a: bool = False
 ) -> Callable[[np.ndarray, bool], tuple[np.ndarray, np.ndarray]]:
     ensure_lapack()
-    _check_linalg_matrix(A, ndim=2, func_name="lu_factor")
+    _check_linalg_matrix(A, ndim=2, dtype=Float, func_name="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)

pytensor/link/numba/dispatch/linalg/solve/lu_solve.py

@@ -87,6 +87,7 @@ def impl(
             NRHS,
             LU.ctypes,
             LDA,
+            # TODO: Does this work with any int dtype?
             IPIV.ctypes,
             B_copy.ctypes,
             LDB,

pytensor/link/numba/dispatch/linalg/solve/norm.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 from numba.core.extending import overload
+from numba.core.types import Float
 from numba.np.linalg import _copy_to_fortran_order, ensure_lapack
 
 from pytensor.link.numba.dispatch.linalg._LAPACK import (
@@ -27,7 +28,7 @@ def xlange_impl(
     largest absolute value of a matrix A.
     """
     ensure_lapack()
-    _check_linalg_matrix(A, ndim=2, func_name="norm")
+    _check_linalg_matrix(A, ndim=2, dtype=Float, func_name="norm")
     dtype = A.dtype
     numba_lange = _LAPACK().numba_xlange(dtype)
 

pytensor/link/numba/dispatch/linalg/solve/tridiagonal.py

@@ -355,6 +355,15 @@ def numba_funcify_LUFactorTridiagonal(op: LUFactorTridiagonal, node, **kwargs):
 
     @numba_basic.numba_njit(cache=False)
     def lu_factor_tridiagonal(dl, d, du):
+        if d.size == 0:
+            return (
+                np.zeros(dl.shape, dtype=out_dtype),
+                np.zeros(d.shape, dtype=out_dtype),
+                np.zeros(du.shape, dtype=out_dtype),
+                np.zeros(d.shape, dtype=out_dtype),
+                np.zeros(d.shape, dtype="int32"),
+            )
+
         if d.dtype != out_dtype:
             d.dtype = out_dtype
         if dl.dtype != out_dtype:
@@ -382,11 +391,18 @@ def numba_funcify_SolveLUFactorTridiagonal(
         return generate_fallback_impl(op, node=node)
     out_dtype = node.outputs[0].type.numpy_dtype
 
+    b_ndim = op.b_ndim
     overwrite_b = op.overwrite_b
     transposed = op.transposed
 
     @numba_basic.numba_njit(cache=False)
     def solve_lu_factor_tridiagonal(dl, d, du, du2, ipiv, b):
+        if d.size == 0:
+            if b_ndim == 1:
+                return np.zeros(d.shape, dtype=out_dtype)
+            else:
+                return np.zeros((d.shape[0], b.shape[1]), dtype=out_dtype)
+
         if dl.dtype != out_dtype:
             dl = dl.astype(out_dtype)
         if d.dtype != out_dtype:

pytensor/link/numba/dispatch/slinalg.py

@@ -74,6 +74,9 @@ def numba_funcify_Cholesky(op, node, **kwargs):
 
     @numba_basic.numba_njit
     def cholesky(a):
+        if a.size == 0:
+            return np.zeros(a.shape, dtype=out_dtype)
+
         if discrete_inp:
             a = a.astype(out_dtype)
         elif check_finite:
@@ -114,7 +117,8 @@ def numba_pivot_to_permutation(piv):
 
         return np.argsort(p_inv)
 
-    return numba_pivot_to_permutation
+    cache_key = 1
+    return numba_pivot_to_permutation, cache_key
 
 
 @numba_funcify.register(LU)
@@ -134,6 +138,18 @@ def numba_funcify_LU(op, node, **kwargs):
 
     @numba_basic.numba_njit
     def lu(a):
+        if a.size == 0:
+            L = np.zeros(a.shape, dtype=a.dtype)
+            U = np.zeros(a.shape, dtype=a.dtype)
+            if permute_l:
+                return L, U
+            elif p_indices:
+                P = np.zeros(a.shape[0], dtype="int32")
+                return P, L, U
+            else:
+                P = np.zeros(a.shape, dtype=a.dtype)
+                return P, L, U
+
         if discrete_inp:
             a = a.astype(out_dtype)
         elif check_finite:
@@ -187,6 +203,12 @@ def numba_funcify_LUFactor(op, node, **kwargs):
 
     @numba_basic.numba_njit
     def lu_factor(a):
+        if a.size == 0:
+            return (
+                np.zeros(a.shape, dtype=out_dtype),
+                np.zeros(a.shape[0], dtype="int32"),
+            )
+
         if discrete_inp:
             a = a.astype(out_dtype)
         elif check_finite:
@@ -226,7 +248,7 @@ def block_diag(*arrs):
 
 @numba_funcify.register(Solve)
 def numba_funcify_Solve(op, node, **kwargs):
-    A_dtype, b_dtype = (i.numpy_dtype for i in node.inputs)
+    A_dtype, b_dtype = (i.type.numpy_dtype for i in node.inputs)
     out_dtype = node.outputs[0].type.numpy_dtype
 
     if A_dtype.kind == "c" or b_dtype.kind == "c":
@@ -269,6 +291,9 @@ def numba_funcify_Solve(op, node, **kwargs):
 
     @numba_basic.numba_njit
     def solve(a, b):
+        if b.size == 0:
+            return np.zeros(b.shape, dtype=out_dtype)
+
         if must_cast_A:
             a = a.astype(out_dtype)
         if must_cast_B:
@@ -297,7 +322,7 @@ def numba_funcify_SolveTriangular(op, node, **kwargs):
     overwrite_b = op.overwrite_b
     b_ndim = op.b_ndim
 
-    A_dtype, b_dtype = (i.numpy_dtype for i in node.inputs)
+    A_dtype, b_dtype = (i.type.numpy_dtype for i in node.inputs)
     out_dtype = node.outputs[0].type.numpy_dtype
 
     if A_dtype.kind == "c" or b_dtype.kind == "c":
@@ -311,6 +336,8 @@ def numba_funcify_SolveTriangular(op, node, **kwargs):
 
     @numba_basic.numba_njit
     def solve_triangular(a, b):
+        if b.size == 0:
+            return np.zeros(b.shape, dtype=out_dtype)
         if must_cast_A:
             a = a.astype(out_dtype)
         if must_cast_B:
@@ -346,7 +373,7 @@ def numba_funcify_CholeskySolve(op, node, **kwargs):
     overwrite_b = op.overwrite_b
     check_finite = op.check_finite
 
-    c_dtype, b_dtype = (i.numpy_dtype for i in node.inputs)
+    c_dtype, b_dtype = (i.type.numpy_dtype for i in node.inputs)
     out_dtype = node.outputs[0].type.numpy_dtype
 
     if c_dtype.kind == "c" or b_dtype.kind == "c":
@@ -360,6 +387,8 @@ def numba_funcify_CholeskySolve(op, node, **kwargs):
 
     @numba_basic.numba_njit
     def cho_solve(c, b):
+        if b.size == 0:
+            return np.zeros(b.shape, dtype=out_dtype)
         if must_cast_c:
             c = c.astype(out_dtype)
         if check_finite:

tests/link/numba/test_slinalg.py

@@ -16,6 +16,13 @@
     LUFactor,
     Solve,
     SolveTriangular,
+    cho_solve,
+    cholesky,
+    lu,
+    lu_factor,
+    lu_solve,
+    solve,
+    solve_triangular,
 )
 from tests.link.numba.test_basic import compare_numba_and_py, numba_inplace_mode
 
@@ -483,6 +490,27 @@ def test_lu_solve(
         # Can never destroy non-contiguous inputs
         np.testing.assert_allclose(b_val_not_contig, b_val)
 
+    @pytest.mark.parametrize(
+        "solve_op",
+        [solve, solve_triangular, cho_solve, lu_solve],
+        ids=lambda x: x.__name__,
+    )
+    def test_empty(self, solve_op):
+        a = pt.matrix("x")
+        b = pt.vector("b")
+        if solve_op is cho_solve:
+            out = solve_op((a, True), b)
+        elif solve_op is lu_solve:
+            out = solve_op((a, b.astype("int32")), b)
+        else:
+            out = solve_op(a, b)
+        compare_numba_and_py(
+            [a, b],
+            [out],
+            [np.zeros((0, 0)), np.zeros(0)],
+            eval_obj_mode=False,  # pivot_to_permutation seems to still be jitted despite the monkey patching
+        )
+
 
 class TestDecompositions:
     @pytest.mark.parametrize("lower", [True, False], ids=lambda x: f"lower={x}")
@@ -750,6 +778,20 @@ def test_qr(self, mode, pivoting, overwrite_a):
         # Cannot destroy non-contiguous input
         np.testing.assert_allclose(val_not_contig, A_val)
 
+    @pytest.mark.parametrize(
+        "decomp_op", (cholesky, lu, lu_factor), ids=lambda x: x.__name__
+    )
+    def test_empty(self, decomp_op):
+        x = pt.matrix("x")
+        outs = decomp_op(x)
+        if not isinstance(outs, tuple | list):
+            outs = [outs]
+        compare_numba_and_py(
+            [x],
+            outs,
+            [np.zeros((0, 0))],
+        )
+
 
 def test_block_diag():
     A = pt.matrix("A")