Numba linalg: handle dtypes more strictly

Author: ricardoV94

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

@@ -1,15 +1,15 @@
 import numpy as np
 from numba.core.extending import overload
 from numba.np.linalg import _copy_to_fortran_order, ensure_lapack
+from numba.types import Float
 from scipy import linalg
 
 from pytensor.link.numba.dispatch.linalg._LAPACK import (
     _LAPACK,
-    _get_underlying_float,
     int_ptr_to_val,
     val_to_int_ptr,
 )
-from pytensor.link.numba.dispatch.linalg.utils import _check_scipy_linalg_matrix
+from pytensor.link.numba.dispatch.linalg.utils import _check_linalg_matrix
 
 
 def _cholesky(a, lower=False, overwrite_a=False, check_finite=True):
@@ -24,9 +24,9 @@ def _cholesky(a, lower=False, overwrite_a=False, check_finite=True):
 @overload(_cholesky)
 def cholesky_impl(A, lower=0, overwrite_a=False, check_finite=True):
     ensure_lapack()
-    _check_scipy_linalg_matrix(A, "cholesky")
+    _check_linalg_matrix(A, ndim=2, dtype=Float, func_name="cholesky")
     dtype = A.dtype
-    w_type = _get_underlying_float(dtype)
+
     numba_potrf = _LAPACK().numba_xpotrf(dtype)
 
     def impl(A, lower=0, overwrite_a=False, check_finite=True):
@@ -47,7 +47,7 @@ def impl(A, lower=0, overwrite_a=False, check_finite=True):
         numba_potrf(
             UPLO,
             N,
-            A_copy.view(w_type).ctypes,
+            A_copy.ctypes,
             LDA,
             INFO,
         )

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

@@ -3,12 +3,13 @@
 
 import numpy as np
 from numba.core.extending import overload
+from numba.core.types import Float
 from numba.np.linalg import ensure_lapack
 from scipy import linalg
 
 from pytensor.link.numba.dispatch import basic as numba_basic
 from pytensor.link.numba.dispatch.linalg.decomposition.lu_factor import _getrf
-from pytensor.link.numba.dispatch.linalg.utils import _check_scipy_linalg_matrix
+from pytensor.link.numba.dispatch.linalg.utils import _check_linalg_matrix
 
 
 @numba_basic.numba_njit
@@ -116,7 +117,7 @@ def lu_impl_1(
     False. Returns a tuple of (perm, L, U), where perm an integer array of row swaps, such that L[perm] @ U = A.
     """
     ensure_lapack()
-    _check_scipy_linalg_matrix(a, "lu")
+    _check_linalg_matrix(a, ndim=2, dtype=Float, func_name="lu")
     dtype = a.dtype
 
     def impl(
@@ -146,7 +147,7 @@ def lu_impl_2(
     """
 
     ensure_lapack()
-    _check_scipy_linalg_matrix(a, "lu")
+    _check_linalg_matrix(a, ndim=2, dtype=Float, func_name="lu")
     dtype = a.dtype
 
     def impl(
@@ -179,7 +180,7 @@ def lu_impl_3(
     False. Returns a tuple of (P, L, U), such that P @ L @ U = A.
     """
     ensure_lapack()
-    _check_scipy_linalg_matrix(a, "lu")
+    _check_linalg_matrix(a, ndim=2, dtype=Float, func_name="lu")
     dtype = a.dtype
 
     def impl(

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

@@ -3,18 +3,16 @@
 
 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
 
 from pytensor.link.numba.dispatch.linalg._LAPACK import (
     _LAPACK,
-    _get_underlying_float,
     int_ptr_to_val,
     val_to_int_ptr,
 )
-from pytensor.link.numba.dispatch.linalg.utils import (
-    _check_scipy_linalg_matrix,
-)
+from pytensor.link.numba.dispatch.linalg.utils import _check_linalg_matrix
 
 
 def _getrf(A, overwrite_a=False) -> tuple[np.ndarray, np.ndarray, int]:
@@ -38,9 +36,8 @@ def getrf_impl(
     A: np.ndarray, overwrite_a: bool = False
 ) -> Callable[[np.ndarray, bool], tuple[np.ndarray, np.ndarray, int]]:
     ensure_lapack()
-    _check_scipy_linalg_matrix(A, "getrf")
+    _check_linalg_matrix(A, ndim=2, dtype=Float, func_name="getrf")
     dtype = A.dtype
-    w_type = _get_underlying_float(dtype)
     numba_getrf = _LAPACK().numba_xgetrf(dtype)
 
     def impl(
@@ -59,7 +56,7 @@ def impl(
         IPIV = np.empty(_N, dtype=np.int32)  # type: ignore
         INFO = val_to_int_ptr(0)
 
-        numba_getrf(M, N, A_copy.view(w_type).ctypes, LDA, IPIV.ctypes, INFO)
+        numba_getrf(M, N, A_copy.ctypes, LDA, IPIV.ctypes, INFO)
 
         return A_copy, IPIV, int_ptr_to_val(INFO)
 
@@ -79,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_scipy_linalg_matrix(A, "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/cholesky.py

@@ -1,17 +1,18 @@
 import numpy as np
 from numba.core.extending import overload
+from numba.core.types import Float
 from numba.np.linalg import ensure_lapack
 from scipy import linalg
 
 from pytensor.link.numba.dispatch.linalg._LAPACK import (
     _LAPACK,
-    _get_underlying_float,
     int_ptr_to_val,
     val_to_int_ptr,
 )
 from pytensor.link.numba.dispatch.linalg.solve.utils import _solve_check_input_shapes
 from pytensor.link.numba.dispatch.linalg.utils import (
-    _check_scipy_linalg_matrix,
+    _check_dtypes_match,
+    _check_linalg_matrix,
     _copy_to_fortran_order_even_if_1d,
     _solve_check,
 )
@@ -31,10 +32,10 @@ def _cho_solve(
 @overload(_cho_solve)
 def cho_solve_impl(C, B, lower=False, overwrite_b=False, check_finite=True):
     ensure_lapack()
-    _check_scipy_linalg_matrix(C, "cho_solve")
-    _check_scipy_linalg_matrix(B, "cho_solve")
+    _check_linalg_matrix(C, ndim=2, dtype=Float, func_name="cho_solve")
+    _check_linalg_matrix(B, ndim=(1, 2), dtype=Float, func_name="cho_solve")
+    _check_dtypes_match((C, B), func_name="cho_solve")
     dtype = C.dtype
-    w_type = _get_underlying_float(dtype)
     numba_potrs = _LAPACK().numba_xpotrs(dtype)
 
     def impl(C, B, lower=False, overwrite_b=False, check_finite=True):
@@ -71,9 +72,9 @@ def impl(C, B, lower=False, overwrite_b=False, check_finite=True):
             UPLO,
             N,
             NRHS,
-            C_f.view(w_type).ctypes,
+            C_f.ctypes,
             LDA,
-            B_copy.view(w_type).ctypes,
+            B_copy.ctypes,
             LDB,
             INFO,
         )

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

@@ -2,12 +2,12 @@
 
 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
 
 from pytensor.link.numba.dispatch.linalg._LAPACK import (
     _LAPACK,
-    _get_underlying_float,
     int_ptr_to_val,
     val_to_int_ptr,
 )
@@ -16,7 +16,8 @@
 from pytensor.link.numba.dispatch.linalg.solve.norm import _xlange
 from pytensor.link.numba.dispatch.linalg.solve.utils import _solve_check_input_shapes
 from pytensor.link.numba.dispatch.linalg.utils import (
-    _check_scipy_linalg_matrix,
+    _check_dtypes_match,
+    _check_linalg_matrix,
     _solve_check,
 )
 
@@ -37,9 +38,8 @@ def xgecon_impl(
     Compute the condition number of a matrix A.
     """
     ensure_lapack()
-    _check_scipy_linalg_matrix(A, "gecon")
+    _check_linalg_matrix(A, ndim=2, dtype=Float, func_name="gecon")
     dtype = A.dtype
-    w_type = _get_underlying_float(dtype)
     numba_gecon = _LAPACK().numba_xgecon(dtype)
 
     def impl(A: np.ndarray, A_norm: float, norm: str) -> tuple[np.ndarray, int]:
@@ -58,11 +58,11 @@ def impl(A: np.ndarray, A_norm: float, norm: str) -> tuple[np.ndarray, int]:
         numba_gecon(
             NORM,
             N,
-            A_copy.view(w_type).ctypes,
+            A_copy.ctypes,
             LDA,
-            A_NORM.view(w_type).ctypes,
-            RCOND.view(w_type).ctypes,
-            WORK.view(w_type).ctypes,
+            A_NORM.ctypes,
+            RCOND.ctypes,
+            WORK.ctypes,
             IWORK.ctypes,
             INFO,
         )
@@ -106,8 +106,9 @@ def solve_gen_impl(
     transposed: bool,
 ) -> Callable[[np.ndarray, np.ndarray, bool, bool, bool, bool, bool], np.ndarray]:
     ensure_lapack()
-    _check_scipy_linalg_matrix(A, "solve")
-    _check_scipy_linalg_matrix(B, "solve")
+    _check_linalg_matrix(A, ndim=2, dtype=Float, func_name="solve")
+    _check_linalg_matrix(B, ndim=(1, 2), dtype=Float, func_name="solve")
+    _check_dtypes_match((A, B), "solve")
 
     def impl(
         A: np.ndarray,

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

@@ -3,18 +3,19 @@
 
 import numpy as np
 from numba.core.extending import overload
+from numba.core.types import Float, int32
 from numba.np.linalg import _copy_to_fortran_order, ensure_lapack
 from scipy import linalg
 
 from pytensor.link.numba.dispatch.linalg._LAPACK import (
     _LAPACK,
-    _get_underlying_float,
     int_ptr_to_val,
     val_to_int_ptr,
 )
 from pytensor.link.numba.dispatch.linalg.solve.utils import _solve_check_input_shapes
 from pytensor.link.numba.dispatch.linalg.utils import (
-    _check_scipy_linalg_matrix,
+    _check_dtypes_match,
+    _check_linalg_matrix,
     _copy_to_fortran_order_even_if_1d,
     _solve_check,
     _trans_char_to_int,
@@ -44,10 +45,11 @@ def getrs_impl(
     [np.ndarray, np.ndarray, np.ndarray, _Trans, bool], tuple[np.ndarray, int]
 ]:
     ensure_lapack()
-    _check_scipy_linalg_matrix(LU, "getrs")
-    _check_scipy_linalg_matrix(B, "getrs")
+    _check_linalg_matrix(LU, ndim=2, dtype=Float, func_name="getrs")
+    _check_linalg_matrix(B, ndim=(1, 2), dtype=Float, func_name="getrs")
+    _check_dtypes_match((LU, B), func_name="getrs")
+    _check_linalg_matrix(IPIV, ndim=1, dtype=int32, func_name="getrs")
     dtype = LU.dtype
-    w_type = _get_underlying_float(dtype)
     numba_getrs = _LAPACK().numba_xgetrs(dtype)
 
     def impl(
@@ -84,10 +86,10 @@ def impl(
             TRANS,
             N,
             NRHS,
-            LU.view(w_type).ctypes,
+            LU.ctypes,
             LDA,
             IPIV.ctypes,
-            B_copy.view(w_type).ctypes,
+            B_copy.ctypes,
             LDB,
             INFO,
         )
@@ -124,8 +126,10 @@ def lu_solve_impl(
     check_finite: bool,
 ) -> Callable[[np.ndarray, np.ndarray, np.ndarray, _Trans, bool, bool], np.ndarray]:
     ensure_lapack()
-    _check_scipy_linalg_matrix(lu_and_piv[0], "lu_solve")
-    _check_scipy_linalg_matrix(b, "lu_solve")
+    lu, _piv = lu_and_piv
+    _check_linalg_matrix(lu, ndim=2, dtype=Float, func_name="lu_solve")
+    _check_linalg_matrix(b, ndim=(1, 2), dtype=Float, func_name="lu_solve")
+    _check_dtypes_match((lu, b), func_name="lu_solve")
 
     def impl(
         lu: np.ndarray,

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

@@ -2,14 +2,14 @@
 
 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 (
     _LAPACK,
-    _get_underlying_float,
     val_to_int_ptr,
 )
-from pytensor.link.numba.dispatch.linalg.utils import _check_scipy_linalg_matrix
+from pytensor.link.numba.dispatch.linalg.utils import _check_linalg_matrix
 
 
 def _xlange(A: np.ndarray, order: str | None = None) -> float:
@@ -28,9 +28,8 @@ def xlange_impl(
     largest absolute value of a matrix A.
     """
     ensure_lapack()
-    _check_scipy_linalg_matrix(A, "norm")
+    _check_linalg_matrix(A, ndim=2, dtype=Float, func_name="norm")
     dtype = A.dtype
-    w_type = _get_underlying_float(dtype)
     numba_lange = _LAPACK().numba_xlange(dtype)
 
     def impl(A: np.ndarray, order: str | None = None):
@@ -49,9 +48,7 @@ def impl(A: np.ndarray, order: str | None = None):
         )
         WORK = np.empty(_M, dtype=dtype)  # type: ignore
 
-        result = numba_lange(
-            NORM, M, N, A_copy.view(w_type).ctypes, LDA, WORK.view(w_type).ctypes
-        )
+        result = numba_lange(NORM, M, N, A_copy.ctypes, LDA, WORK.ctypes)
 
         return result