Use more specialized scipy sparse imports

Author: ricardoV94

pytensor/sparse/basic.py

@@ -14,6 +14,7 @@
 import numpy as np
 import scipy.sparse
 from numpy.lib.stride_tricks import as_strided
+from scipy.sparse import issparse, spmatrix
 
 import pytensor
 from pytensor import _as_symbolic, as_symbolic
@@ -70,20 +71,6 @@
 )
 
 
-sparse_formats = ["csc", "csr"]
-
-"""
-Types of sparse matrices to use for testing.
-
-"""
-_mtypes = [scipy.sparse.csc_matrix, scipy.sparse.csr_matrix]
-# _mtypes = [sparse.csc_matrix, sparse.csr_matrix, sparse.dok_matrix,
-# sparse.lil_matrix, sparse.coo_matrix]
-# * new class ``dia_matrix`` : the sparse DIAgonal format
-# * new class ``bsr_matrix`` : the Block CSR format
-_mtype_to_str = {scipy.sparse.csc_matrix: "csc", scipy.sparse.csr_matrix: "csr"}
-
-
 def _is_sparse_variable(x):
     """
 
@@ -134,7 +121,7 @@ def _is_dense(x):
         L{numpy.ndarray}).
 
     """
-    if not isinstance(x, scipy.sparse.spmatrix | np.ndarray):
+    if not isinstance(x, spmatrix | np.ndarray):
         raise NotImplementedError(
             "this function should only be called on "
             "sparse.scipy.sparse.spmatrix or "
@@ -144,7 +131,7 @@ def _is_dense(x):
     return isinstance(x, np.ndarray)
 
 
-@_as_symbolic.register(scipy.sparse.spmatrix)
+@_as_symbolic.register(spmatrix)
 def as_symbolic_sparse(x, **kwargs):
     return as_sparse_variable(x, **kwargs)
 
@@ -198,7 +185,7 @@ def as_sparse_variable(x, name=None, ndim=None, **kwargs):
 
 
 def constant(x, name=None):
-    if not isinstance(x, scipy.sparse.spmatrix):
+    if not isinstance(x, spmatrix):
         raise TypeError("sparse.constant must be called on a scipy.sparse.spmatrix")
     try:
         return SparseConstant(
@@ -3337,7 +3324,7 @@ def perform(self, node, inp, out_):
         x, y = inp
         (out,) = out_
         rval = x.dot(y)
-        if not scipy.sparse.issparse(rval):
+        if not issparse(rval):
             rval = getattr(scipy.sparse, x.format + "_matrix")(rval)
         # x.dot call tocsr() that will "upcast" to ['int8', 'uint8', 'short',
         # 'ushort', 'intc', 'uintc', 'longlong', 'ulonglong', 'single',
@@ -3604,7 +3591,7 @@ def perform(self, node, inputs, outputs):
                 # the following dot product can result in a scalar or
                 # a (1, 1) sparse matrix.
                 dot_val = np.dot(g_ab[i], b[j].T)
-                if isinstance(dot_val, scipy.sparse.spmatrix):
+                if isinstance(dot_val, spmatrix):
                     dot_val = dot_val[0, 0]
                 g_a_data[i_idx] = dot_val
         out[0] = g_a_data
@@ -3738,7 +3725,7 @@ def perform(self, node, inputs, outputs):
                 # the following dot product can result in a scalar or
                 # a (1, 1) sparse matrix.
                 dot_val = np.dot(g_ab[i], b[j].T)
-                if isinstance(dot_val, scipy.sparse.spmatrix):
+                if isinstance(dot_val, spmatrix):
                     dot_val = dot_val[0, 0]
                 g_a_data[j_idx] = dot_val
         out[0] = g_a_data
@@ -3955,9 +3942,9 @@ def make_node(self, x, y):
         # Sparse dot product should have at least one sparse variable
         # as input. If the other one is not sparse, it has to be converted
         # into a tensor.
-        if isinstance(x, scipy.sparse.spmatrix):
+        if isinstance(x, spmatrix):
             x = as_sparse_variable(x)
-        if isinstance(y, scipy.sparse.spmatrix):
+        if isinstance(y, spmatrix):
             y = as_sparse_variable(y)
 
         x_is_sparse_var = _is_sparse_variable(x)
@@ -4147,7 +4134,7 @@ def perform(self, node, inputs, outputs):
             raise TypeError(x)
 
         rval = x * y
-        if isinstance(rval, scipy.sparse.spmatrix):
+        if isinstance(rval, spmatrix):
             rval = rval.toarray()
         if rval.dtype == alpha.dtype:
             rval *= alpha  # Faster because operation is inplace

pytensor/sparse/rewriting.py

@@ -1,5 +1,5 @@
 import numpy as np
-import scipy
+from scipy.sparse import csc_matrix, csr_matrix
 
 import pytensor
 import pytensor.scalar as ps
@@ -279,9 +279,7 @@ def make_node(self, a_val, a_ind, a_ptr, a_nrows, b):
     def perform(self, node, inputs, outputs):
         (a_val, a_ind, a_ptr, a_nrows, b) = inputs
         (out,) = outputs
-        a = scipy.sparse.csc_matrix(
-            (a_val, a_ind, a_ptr), (a_nrows, b.shape[0]), copy=False
-        )
+        a = csc_matrix((a_val, a_ind, a_ptr), (a_nrows, b.shape[0]), copy=False)
         # out[0] = a.dot(b)
         out[0] = np.asarray(a * b, dtype=node.outputs[0].type.dtype)
         assert _is_dense(out[0])  # scipy 0.7 automatically converts to dense
@@ -478,7 +476,7 @@ def make_node(self, a_val, a_ind, a_ptr, b):
     def perform(self, node, inputs, outputs):
         (a_val, a_ind, a_ptr, b) = inputs
         (out,) = outputs
-        a = scipy.sparse.csr_matrix(
+        a = csr_matrix(
             (a_val, a_ind, a_ptr), (len(a_ptr) - 1, b.shape[0]), copy=True
         )  # use view_map before setting this to False
         # out[0] = a.dot(b)

pytensor/sparse/sharedvar.py

@@ -1,6 +1,6 @@
 import copy
 
-import scipy.sparse
+from scipy.sparse import spmatrix
 
 from pytensor.compile import shared_constructor
 from pytensor.sparse.basic import SparseTensorType, SparseVariable
@@ -13,7 +13,7 @@ def format(self):
         return self.type.format
 
 
-@shared_constructor.register(scipy.sparse.spmatrix)
+@shared_constructor.register(spmatrix)
 def sparse_constructor(
     value, name=None, strict=False, allow_downcast=None, borrow=False, format=None
 ):

pytensor/sparse/type.py

@@ -2,7 +2,7 @@
 from typing import Literal
 
 import numpy as np
-import scipy.sparse
+from scipy.sparse import bsr_matrix, csc_matrix, csr_matrix, issparse, spmatrix
 
 import pytensor
 from pytensor import scalar as ps
@@ -23,14 +23,14 @@ def _is_sparse(x):
         True iff x is a L{scipy.sparse.spmatrix} (and not a L{numpy.ndarray}).
 
     """
-    if not isinstance(x, scipy.sparse.spmatrix | np.ndarray | tuple | list):
+    if not isinstance(x, spmatrix | np.ndarray | tuple | list):
         raise NotImplementedError(
             "this function should only be called on "
             "sparse.scipy.sparse.spmatrix or "
             "numpy.ndarray, not,",
             x,
         )
-    return isinstance(x, scipy.sparse.spmatrix)
+    return isinstance(x, spmatrix)
 
 
 class SparseTensorType(TensorType, HasDataType):
@@ -44,9 +44,9 @@ class SparseTensorType(TensorType, HasDataType):
 
     __props__ = ("dtype", "format", "shape")
     format_cls = {
-        "csr": scipy.sparse.csr_matrix,
-        "csc": scipy.sparse.csc_matrix,
-        "bsr": scipy.sparse.bsr_matrix,
+        "csr": csr_matrix,
+        "csc": csc_matrix,
+        "bsr": bsr_matrix,
     }
     dtype_specs_map = {
         "float32": (float, "npy_float32", "NPY_FLOAT32"),
@@ -187,7 +187,7 @@ def values_eq_approx(self, a, b, eps=1e-6):
         # WARNING: equality comparison of sparse matrices is not fast or easy
         # we definitely do not want to be doing this un-necessarily during
         # a FAST_RUN computation..
-        if not (scipy.sparse.issparse(a) and scipy.sparse.issparse(b)):
+        if not (issparse(a) and issparse(b)):
             return False
         diff = abs(a - b)
         if diff.nnz == 0:
@@ -203,14 +203,10 @@ def values_eq(self, a, b):
         # WARNING: equality comparison of sparse matrices is not fast or easy
         # we definitely do not want to be doing this un-necessarily during
         # a FAST_RUN computation..
-        return (
-            scipy.sparse.issparse(a)
-            and scipy.sparse.issparse(b)
-            and abs(a - b).sum() == 0.0
-        )
+        return issparse(a) and issparse(b) and abs(a - b).sum() == 0.0
 
     def is_valid_value(self, a):
-        return scipy.sparse.issparse(a) and (a.format == self.format)
+        return issparse(a) and (a.format == self.format)
 
     def get_shape_info(self, obj):
         obj = self.filter(obj)