Use get_blas_funcs instead of storing fortran objects in a dict

Author: lucianopaz

pytensor/tensor/blas.py

@@ -110,22 +110,9 @@
 _logger = logging.getLogger("pytensor.tensor.blas")
 
 try:
-    import scipy.linalg.blas
+    from scipy.linalg.blas import get_blas_funcs
 
     have_fblas = True
-    try:
-        fblas = scipy.linalg.blas.fblas
-    except AttributeError:
-        # A change merged in Scipy development version on 2012-12-02 replaced
-        # `scipy.linalg.blas.fblas` with `scipy.linalg.blas`.
-        # See http://github.com/scipy/scipy/pull/358
-        fblas = scipy.linalg.blas
-    _blas_gemv_fns = {
-        np.dtype("float32"): fblas.sgemv,
-        np.dtype("float64"): fblas.dgemv,
-        np.dtype("complex64"): fblas.cgemv,
-        np.dtype("complex128"): fblas.zgemv,
-    }
 except ImportError as e:
     have_fblas = False
     # This is used in Gemv and ScipyGer. We use CGemv and CGer
@@ -146,18 +133,18 @@ def check_init_y():
     if check_init_y._result is None:
         if not have_fblas:
             check_init_y._result = False
-
-        y = float("NaN") * np.ones((2,))
-        x = np.ones((2,))
-        A = np.ones((2, 2))
-        gemv = _blas_gemv_fns[y.dtype]
-        gemv(1.0, A.T, x, 0.0, y, overwrite_y=True, trans=True)
-        check_init_y._result = np.isnan(y).any()
+        else:
+            y = float("NaN") * np.ones((2,))
+            x = np.ones((2,))
+            A = np.ones((2, 2))
+            gemv = get_blas_funcs(names="gemv", dtype=y.dtype)
+            gemv(1.0, A.T, x, 0.0, y, overwrite_y=True, trans=True)
+            check_init_y._result = np.isnan(y).any()
 
     return check_init_y._result
 
 
-check_init_y._result = None
+check_init_y._result = None  # type: ignore
 
 
 class Gemv(Op):
@@ -210,14 +197,11 @@ def make_node(self, y, alpha, A, x, beta):
 
     def perform(self, node, inputs, out_storage):
         y, alpha, A, x, beta = inputs
-        if (
-            have_fblas
-            and y.shape[0] != 0
-            and x.shape[0] != 0
-            and y.dtype in _blas_gemv_fns
-        ):
-            gemv = _blas_gemv_fns[y.dtype]
-
+        try:
+            gemv = get_blas_funcs(names="gemv", dtype=y.dtype)
+        except Exception:
+            gemv = None
+        if have_fblas and y.shape[0] != 0 and x.shape[0] != 0 and gemv is not None:
             if A.shape[0] != y.shape[0] or A.shape[1] != x.shape[0]:
                 raise ValueError(
                     "Incompatible shapes for gemv "

pytensor/tensor/blas_scipy.py

@@ -2,34 +2,22 @@
 Implementations of BLAS Ops based on scipy's BLAS bindings.
 """
 
-import numpy as np
 
 from pytensor.tensor.blas import Ger, have_fblas
 
 
 if have_fblas:
-    from pytensor.tensor.blas import fblas
-
-    _blas_ger_fns = {
-        np.dtype("float32"): fblas.sger,
-        np.dtype("float64"): fblas.dger,
-        np.dtype("complex64"): fblas.cgeru,
-        np.dtype("complex128"): fblas.zgeru,
-    }
+    from scipy.linalg.blas import get_blas_funcs
 
 
 class ScipyGer(Ger):
-    def prepare_node(self, node, storage_map, compute_map, impl):
-        if impl == "py":
-            node.tag.local_ger = _blas_ger_fns[np.dtype(node.inputs[0].type.dtype)]
-
     def perform(self, node, inputs, output_storage):
         cA, calpha, cx, cy = inputs
         (cZ,) = output_storage
         # N.B. some versions of scipy (e.g. mine) don't actually work
         # in-place on a, even when I tell it to.
         A = cA
-        local_ger = node.tag.local_ger
+        local_ger = get_blas_funcs(names="ger", dtype=cA.dtype)
         if A.size == 0:
             # We don't have to compute anything, A is empty.
             # We need this special case because Numpy considers it

scripts/mypy-failing.txt

@@ -17,7 +17,6 @@ pytensor/scalar/basic.py
 pytensor/sparse/basic.py
 pytensor/sparse/type.py
 pytensor/tensor/basic.py
-pytensor/tensor/blas.py
 pytensor/tensor/blas_c.py
 pytensor/tensor/blas_headers.py
 pytensor/tensor/elemwise.py
@@ -31,4 +30,4 @@ pytensor/tensor/slinalg.py
 pytensor/tensor/subtensor.py
 pytensor/tensor/type.py
 pytensor/tensor/type_other.py
-pytensor/tensor/variable.py
+pytensor/tensor/variable.py