All hail mypy

Author: jessegrabowski

pytensor/link/numba/dispatch/linalg/dot/banded.py

@@ -56,10 +56,14 @@ def _gbmv(
     A_banded = A_to_banded(A, kl=kl, ku=ku)
 
     incx = x.strides[0] // x.itemsize
-    incy = y.strides[0] // y.itemsize if y is not None else 1
-
     offx = 0 if incx >= 0 else -x.size + 1
-    offy = 0 if incy >= 0 else -y.size + 1
+
+    if y is not None:
+        incy = y.strides[0] // y.itemsize
+        offy = 0 if incy >= 0 else -y.size + 1
+    else:
+        incy = 1
+        offy = 0
 
     return fn(
         m=m,

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

@@ -1,4 +1,5 @@
 from collections.abc import Callable
+from typing import cast as type_cast
 
 import numpy as np
 from numba.core.extending import overload
@@ -33,12 +34,16 @@ def _matrix_vector_product(
     (fn,) = linalg.get_blas_funcs(("gemv",), (A, x))
 
     incx = x.strides[0] // x.itemsize
-    incy = y.strides[0] // y.itemsize if y is not None else 1
-
     offx = 0 if incx >= 0 else -x.size + 1
-    offy = 0 if incy >= 0 else -y.size + 1
 
-    return fn(
+    if y is not None:
+        incy = y.strides[0] // y.itemsize if y is not None else 1
+        offy = 0 if incy >= 0 else -y.size + 1
+    else:
+        incy = 1
+        offy = 0
+
+    res = fn(
         alpha=alpha,
         a=A,
         x=x,
@@ -52,6 +57,8 @@ def _matrix_vector_product(
         trans=trans,
     )
 
+    return type_cast(np.ndarray, res)
+
 
 @overload(_matrix_vector_product)
 def matrix_vector_product_impl(
@@ -63,7 +70,15 @@ def matrix_vector_product_impl(
     overwrite_y: bool = False,
     trans: int = 1,
 ) -> Callable[
-    [float, np.ndarray, np.ndarray, float, np.ndarray, int, int, int, int, int],
+    [
+        np.ndarray,
+        np.ndarray,
+        np.ndarray,
+        np.ndarray | None,
+        np.ndarray | None,
+        bool,
+        int,
+    ],
     np.ndarray,
 ]:
     ensure_lapack()

pytensor/tensor/blas.py

@@ -1941,12 +1941,23 @@ def banded_gemv(
     out: Tensor
         The matrix multiplication result
     """
+    A = as_tensor_variable(A)
+    x = as_tensor_variable(x)
+
     if alpha is None:
         alpha = pt.ones((), dtype=A.type.dtype)
+    else:
+        alpha = as_tensor_variable(alpha)
+
     if beta is None:
         beta = pt.zeros((), dtype=A.type.dtype)
+    else:
+        beta = as_tensor_variable(beta)
+
     if y is None:
         y = pt.empty(A.shape[:-1], dtype=A.type.dtype)
+    else:
+        y = as_tensor_variable(y)
 
     return Blockwise(BandedGEMV(lower_diags, upper_diags, overwrite_y=False))(
         A, x, y, alpha, beta