diff --git a/pytensor/compile/compilelock.py b/pytensor/compile/compilelock.py index 83bf42866d..a1697e43d1 100644 --- a/pytensor/compile/compilelock.py +++ b/pytensor/compile/compilelock.py @@ -8,8 +8,6 @@ from contextlib import contextmanager from pathlib import Path -import filelock - from pytensor.configdefaults import config @@ -35,8 +33,9 @@ def force_unlock(lock_dir: os.PathLike): lock_dir : os.PathLike Path to a directory that was locked with `lock_ctx`. """ + from filelock import FileLock - fl = filelock.FileLock(Path(lock_dir) / ".lock") + fl = FileLock(Path(lock_dir) / ".lock") fl.release(force=True) dir_key = f"{lock_dir}-{os.getpid()}" @@ -62,6 +61,8 @@ def lock_ctx( Timeout in seconds for waiting in lock acquisition. Defaults to `pytensor.config.compile__timeout`. """ + from filelock import FileLock + if lock_dir is None: lock_dir = config.compiledir @@ -73,7 +74,7 @@ def lock_ctx( if dir_key not in local_mem._locks: local_mem._locks[dir_key] = True - fl = filelock.FileLock(Path(lock_dir) / ".lock") + fl = FileLock(Path(lock_dir) / ".lock") fl.acquire(timeout=timeout) try: yield diff --git a/pytensor/configdefaults.py b/pytensor/configdefaults.py index a81fd63905..fcc36f0c6f 100644 --- a/pytensor/configdefaults.py +++ b/pytensor/configdefaults.py @@ -3,11 +3,10 @@ import os import platform import re -import shutil -import socket import sys import textwrap from pathlib import Path +from shutil import which import numpy as np @@ -349,7 +348,7 @@ def add_compile_configvars(): # Try to find the full compiler path from the name if param != "": - newp = shutil.which(param) + newp = which(param) if newp is not None: param = newp del newp @@ -1190,7 +1189,7 @@ def _get_home_dir() -> Path: "pytensor_version": pytensor.__version__, "numpy_version": np.__version__, "gxx_version": "xxx", - "hostname": socket.gethostname(), + "hostname": platform.node(), } diff --git a/pytensor/configparser.py b/pytensor/configparser.py index 8c6da4a144..4f71e85240 100644 --- a/pytensor/configparser.py +++ b/pytensor/configparser.py @@ -1,6 +1,5 @@ import logging import os -import shlex import sys import warnings from collections.abc import Callable, Sequence @@ -14,6 +13,7 @@ from functools import wraps from io import StringIO from pathlib import Path +from shlex import shlex from pytensor.utils import hash_from_code @@ -541,7 +541,7 @@ def parse_config_string( Parses a config string (comma-separated key=value components) into a dict. """ config_dict = {} - my_splitter = shlex.shlex(config_string, posix=True) + my_splitter = shlex(config_string, posix=True) my_splitter.whitespace = "," my_splitter.whitespace_split = True for kv_pair in my_splitter: diff --git a/pytensor/d3viz/formatting.py b/pytensor/d3viz/formatting.py index b9fb8ee5a5..df39335c19 100644 --- a/pytensor/d3viz/formatting.py +++ b/pytensor/d3viz/formatting.py @@ -12,13 +12,7 @@ from pytensor.compile import Function, builders from pytensor.graph.basic import Apply, Constant, Variable, graph_inputs from pytensor.graph.fg import FunctionGraph -from pytensor.printing import pydot_imported, pydot_imported_msg - - -try: - from pytensor.printing import pd -except ImportError: - pass +from pytensor.printing import _try_pydot_import class PyDotFormatter: @@ -41,8 +35,7 @@ class PyDotFormatter: def __init__(self, compact=True): """Construct PyDotFormatter object.""" - if not pydot_imported: - raise ImportError("Failed to import pydot. " + pydot_imported_msg) + _try_pydot_import() self.compact = compact self.node_colors = { @@ -115,6 +108,8 @@ def __call__(self, fct, graph=None): pydot.Dot Pydot graph of `fct` """ + pd = _try_pydot_import() + if graph is None: graph = pd.Dot() @@ -356,6 +351,8 @@ def type_to_str(t): def dict_to_pdnode(d): """Create pydot node from dict.""" + pd = _try_pydot_import() + e = dict() for k, v in d.items(): if v is not None: diff --git a/pytensor/graph/rewriting/basic.py b/pytensor/graph/rewriting/basic.py index faec736c98..344d6a1940 100644 --- a/pytensor/graph/rewriting/basic.py +++ b/pytensor/graph/rewriting/basic.py @@ -5,7 +5,6 @@ import functools import inspect import logging -import pdb import sys import time import traceback @@ -237,6 +236,8 @@ def warn(cls, exc, self, rewriter): if config.on_opt_error == "raise": raise exc elif config.on_opt_error == "pdb": + import pdb + pdb.post_mortem(sys.exc_info()[2]) def __init__(self, *rewrites, failure_callback=None): @@ -1752,6 +1753,8 @@ def warn(cls, exc, nav, repl_pairs, node_rewriter, node): _logger.error("TRACEBACK:") _logger.error(traceback.format_exc()) if config.on_opt_error == "pdb": + import pdb + pdb.post_mortem(sys.exc_info()[2]) elif isinstance(exc, AssertionError) or config.on_opt_error == "raise": # We always crash on AssertionError because something may be diff --git a/pytensor/link/c/cmodule.py b/pytensor/link/c/cmodule.py index f1f098edbf..c992d0506e 100644 --- a/pytensor/link/c/cmodule.py +++ b/pytensor/link/c/cmodule.py @@ -26,19 +26,12 @@ from typing import TYPE_CHECKING, Protocol, cast import numpy as np -from setuptools._distutils.sysconfig import ( - get_config_h_filename, - get_config_var, - get_python_inc, - get_python_lib, -) # we will abuse the lockfile mechanism when reading and writing the registry from pytensor.compile.compilelock import lock_ctx from pytensor.configdefaults import config, gcc_version_str from pytensor.configparser import BoolParam, StrParam from pytensor.graph.op import Op -from pytensor.link.c.exceptions import CompileError, MissingGXX from pytensor.utils import ( LOCAL_BITWIDTH, flatten, @@ -266,6 +259,8 @@ def list_code(self, ofile=sys.stdout): def _get_ext_suffix(): """Get the suffix for compiled extensions""" + from setuptools._distutils.sysconfig import get_config_var + dist_suffix = get_config_var("EXT_SUFFIX") if dist_suffix is None: dist_suffix = get_config_var("SO") @@ -1697,6 +1692,8 @@ def get_gcc_shared_library_arg(): def std_include_dirs(): + from setuptools._distutils.sysconfig import get_python_inc + numpy_inc_dirs = [np.get_include()] py_inc = get_python_inc() py_plat_spec_inc = get_python_inc(plat_specific=True) @@ -1709,6 +1706,12 @@ def std_include_dirs(): @is_StdLibDirsAndLibsType def std_lib_dirs_and_libs() -> tuple[list[str], ...] | None: + from setuptools._distutils.sysconfig import ( + get_config_var, + get_python_inc, + get_python_lib, + ) + # We cache the results as on Windows, this trigger file access and # this method is called many times. if std_lib_dirs_and_libs.data is not None: @@ -2388,23 +2391,6 @@ def join_options(init_part): # xcode's version. cxxflags.append("-ld64") - if sys.platform == "win32": - # Workaround for https://github.com/Theano/Theano/issues/4926. - # https://github.com/python/cpython/pull/11283/ removed the "hypot" - # redefinition for recent CPython versions (>=2.7.16 and >=3.7.3). - # The following nullifies that redefinition, if it is found. - python_version = sys.version_info[:3] - if (3,) <= python_version < (3, 7, 3): - config_h_filename = get_config_h_filename() - try: - with open(config_h_filename) as config_h: - if any( - line.startswith("#define hypot _hypot") for line in config_h - ): - cxxflags.append("-D_hypot=hypot") - except OSError: - pass - return cxxflags @classmethod @@ -2555,8 +2541,9 @@ def compile_str( """ # TODO: Do not do the dlimport in this function - if not config.cxx: + from pytensor.link.c.exceptions import MissingGXX + raise MissingGXX("g++ not available! We can't compile c code.") if include_dirs is None: @@ -2586,6 +2573,8 @@ def compile_str( cppfile.write("\n") if platform.python_implementation() == "PyPy": + from setuptools._distutils.sysconfig import get_config_var + suffix = "." + get_lib_extension() dist_suffix = get_config_var("SO") @@ -2642,6 +2631,8 @@ def print_command_line_error(): status = p_out[2] if status: + from pytensor.link.c.exceptions import CompileError + tf = tempfile.NamedTemporaryFile( mode="w", prefix="pytensor_compilation_error_", delete=False ) diff --git a/pytensor/link/vm.py b/pytensor/link/vm.py index a9d625a8da..af44af3254 100644 --- a/pytensor/link/vm.py +++ b/pytensor/link/vm.py @@ -19,7 +19,6 @@ from pytensor.configdefaults import config from pytensor.graph.basic import Apply, Constant, Variable from pytensor.link.basic import Container, LocalLinker -from pytensor.link.c.exceptions import MissingGXX from pytensor.link.utils import ( gc_helper, get_destroy_dependencies, @@ -1006,6 +1005,8 @@ def make_vm( compute_map, updated_vars, ): + from pytensor.link.c.exceptions import MissingGXX + pre_call_clear = [storage_map[v] for v in self.no_recycling] try: diff --git a/pytensor/printing.py b/pytensor/printing.py index 9a34317c40..6a18f6e8e5 100644 --- a/pytensor/printing.py +++ b/pytensor/printing.py @@ -26,39 +26,6 @@ IDTypesType = Literal["id", "int", "CHAR", "auto", ""] -pydot_imported = False -pydot_imported_msg = "" -try: - # pydot-ng is a fork of pydot that is better maintained - import pydot_ng as pd - - if pd.find_graphviz(): - pydot_imported = True - else: - pydot_imported_msg = "pydot-ng can't find graphviz. Install graphviz." -except ImportError: - try: - # fall back on pydot if necessary - import pydot as pd - - if hasattr(pd, "find_graphviz"): - if pd.find_graphviz(): - pydot_imported = True - else: - pydot_imported_msg = "pydot can't find graphviz" - else: - pd.Dot.create(pd.Dot()) - pydot_imported = True - except ImportError: - # tests should not fail on optional dependency - pydot_imported_msg = ( - "Install the python package pydot or pydot-ng. Install graphviz." - ) - except Exception as e: - pydot_imported_msg = "An error happened while importing/trying pydot: " - pydot_imported_msg += str(e.args) - - _logger = logging.getLogger("pytensor.printing") VALID_ASSOC = {"left", "right", "either"} @@ -1196,6 +1163,48 @@ def __call__(self, *args): } +def _try_pydot_import(): + pydot_imported = False + pydot_imported_msg = "" + try: + # pydot-ng is a fork of pydot that is better maintained + import pydot_ng as pd + + if pd.find_graphviz(): + pydot_imported = True + else: + pydot_imported_msg = "pydot-ng can't find graphviz. Install graphviz." + except ImportError: + try: + # fall back on pydot if necessary + import pydot as pd + + if hasattr(pd, "find_graphviz"): + if pd.find_graphviz(): + pydot_imported = True + else: + pydot_imported_msg = "pydot can't find graphviz" + else: + pd.Dot.create(pd.Dot()) + pydot_imported = True + except ImportError: + # tests should not fail on optional dependency + pydot_imported_msg = ( + "Install the python package pydot or pydot-ng. Install graphviz." + ) + except Exception as e: + pydot_imported_msg = "An error happened while importing/trying pydot: " + pydot_imported_msg += str(e.args) + + if not pydot_imported: + raise ImportError( + "Failed to import pydot. You must install graphviz " + "and either pydot or pydot-ng for " + f"`pydotprint` to work:\n {pydot_imported_msg}", + ) + return pd + + def pydotprint( fct, outfile: Path | str | None = None, @@ -1288,6 +1297,8 @@ def pydotprint( scan separately after the top level debugprint output. """ + pd = _try_pydot_import() + from pytensor.scan.op import Scan if colorCodes is None: @@ -1320,12 +1331,6 @@ def pydotprint( outputs = fct.outputs topo = fct.toposort() fgraph = fct - if not pydot_imported: - raise RuntimeError( - "Failed to import pydot. You must install graphviz " - "and either pydot or pydot-ng for " - f"`pydotprint` to work:\n {pydot_imported_msg}", - ) g = pd.Dot() diff --git a/pytensor/scalar/math.py b/pytensor/scalar/math.py index a5512c6564..e50e28ac50 100644 --- a/pytensor/scalar/math.py +++ b/pytensor/scalar/math.py @@ -9,8 +9,6 @@ from textwrap import dedent import numpy as np -import scipy.special -import scipy.stats from pytensor.configdefaults import config from pytensor.gradient import grad_not_implemented, grad_undefined @@ -54,7 +52,9 @@ class Erf(UnaryScalarOp): nfunc_spec = ("scipy.special.erf", 1, 1) def impl(self, x): - return scipy.special.erf(x) + from scipy.special import erf + + return erf(x) def L_op(self, inputs, outputs, grads): (x,) = inputs @@ -88,7 +88,9 @@ class Erfc(UnaryScalarOp): nfunc_spec = ("scipy.special.erfc", 1, 1) def impl(self, x): - return scipy.special.erfc(x) + from scipy.special import erfc + + return erfc(x) def L_op(self, inputs, outputs, grads): (x,) = inputs @@ -137,7 +139,9 @@ class Erfcx(UnaryScalarOp): nfunc_spec = ("scipy.special.erfcx", 1, 1) def impl(self, x): - return scipy.special.erfcx(x) + from scipy.special import erfcx + + return erfcx(x) def L_op(self, inputs, outputs, grads): (x,) = inputs @@ -193,7 +197,9 @@ class Erfinv(UnaryScalarOp): nfunc_spec = ("scipy.special.erfinv", 1, 1) def impl(self, x): - return scipy.special.erfinv(x) + from scipy.special import erfinv + + return erfinv(x) def L_op(self, inputs, outputs, grads): (x,) = inputs @@ -228,7 +234,9 @@ class Erfcinv(UnaryScalarOp): nfunc_spec = ("scipy.special.erfcinv", 1, 1) def impl(self, x): - return scipy.special.erfcinv(x) + from scipy.special import erfcinv + + return erfcinv(x) def L_op(self, inputs, outputs, grads): (x,) = inputs @@ -264,7 +272,9 @@ class Owens_t(BinaryScalarOp): @staticmethod def st_impl(h, a): - return scipy.special.owens_t(h, a) + from scipy.special import owens_t + + return owens_t(h, a) def impl(self, h, a): return Owens_t.st_impl(h, a) @@ -293,7 +303,9 @@ class Gamma(UnaryScalarOp): @staticmethod def st_impl(x): - return scipy.special.gamma(x) + from scipy.special import gamma + + return gamma(x) def impl(self, x): return Gamma.st_impl(x) @@ -332,7 +344,9 @@ class GammaLn(UnaryScalarOp): @staticmethod def st_impl(x): - return scipy.special.gammaln(x) + from scipy.special import gammaln + + return gammaln(x) def impl(self, x): return GammaLn.st_impl(x) @@ -376,7 +390,9 @@ class Psi(UnaryScalarOp): @staticmethod def st_impl(x): - return scipy.special.psi(x) + from scipy.special import psi + + return psi(x) def impl(self, x): return Psi.st_impl(x) @@ -467,7 +483,9 @@ class TriGamma(UnaryScalarOp): @staticmethod def st_impl(x): - return scipy.special.polygamma(1, x) + from scipy.special import polygamma + + return polygamma(1, x) def impl(self, x): return TriGamma.st_impl(x) @@ -570,7 +588,9 @@ def output_types_preference(n_type, x_type): @staticmethod def st_impl(n, x): - return scipy.special.polygamma(n, x) + from scipy.special import polygamma + + return polygamma(n, x) def impl(self, n, x): return PolyGamma.st_impl(n, x) @@ -602,7 +622,9 @@ class Chi2SF(BinaryScalarOp): @staticmethod def st_impl(x, k): - return scipy.stats.chi2.sf(x, k) + from scipy.stats import chi2 + + return chi2.sf(x, k) def impl(self, x, k): return Chi2SF.st_impl(x, k) @@ -645,7 +667,9 @@ class GammaInc(BinaryScalarOp): @staticmethod def st_impl(k, x): - return scipy.special.gammainc(k, x) + from scipy.special import gammainc + + return gammainc(k, x) def impl(self, k, x): return GammaInc.st_impl(k, x) @@ -696,7 +720,9 @@ class GammaIncC(BinaryScalarOp): @staticmethod def st_impl(k, x): - return scipy.special.gammaincc(k, x) + from scipy.special import gammaincc + + return gammaincc(k, x) def impl(self, k, x): return GammaIncC.st_impl(k, x) @@ -747,7 +773,9 @@ class GammaIncInv(BinaryScalarOp): @staticmethod def st_impl(k, x): - return scipy.special.gammaincinv(k, x) + from scipy.special import gammaincinv + + return gammaincinv(k, x) def impl(self, k, x): return GammaIncInv.st_impl(k, x) @@ -776,7 +804,9 @@ class GammaIncCInv(BinaryScalarOp): @staticmethod def st_impl(k, x): - return scipy.special.gammainccinv(k, x) + from scipy.special import gammainccinv + + return gammainccinv(k, x) def impl(self, k, x): return GammaIncCInv.st_impl(k, x) @@ -1015,7 +1045,9 @@ class GammaU(BinaryScalarOp): @staticmethod def st_impl(k, x): - return scipy.special.gammaincc(k, x) * scipy.special.gamma(k) + from scipy.special import gamma, gammaincc + + return gammaincc(k, x) * gamma(k) def impl(self, k, x): return GammaU.st_impl(k, x) @@ -1051,7 +1083,9 @@ class GammaL(BinaryScalarOp): @staticmethod def st_impl(k, x): - return scipy.special.gammainc(k, x) * scipy.special.gamma(k) + from scipy.special import gamma, gammainc + + return gammainc(k, x) * gamma(k) def impl(self, k, x): return GammaL.st_impl(k, x) @@ -1087,7 +1121,9 @@ class Jv(BinaryScalarOp): @staticmethod def st_impl(v, x): - return scipy.special.jv(v, x) + from scipy.special import jv + + return jv(v, x) def impl(self, v, x): return self.st_impl(v, x) @@ -1116,7 +1152,9 @@ class J1(UnaryScalarOp): @staticmethod def st_impl(x): - return scipy.special.j1(x) + from scipy.special import j1 + + return j1(x) def impl(self, x): return self.st_impl(x) @@ -1147,7 +1185,9 @@ class J0(UnaryScalarOp): @staticmethod def st_impl(x): - return scipy.special.j0(x) + from scipy.special import j0 + + return j0(x) def impl(self, x): return self.st_impl(x) @@ -1178,7 +1218,9 @@ class Iv(BinaryScalarOp): @staticmethod def st_impl(v, x): - return scipy.special.iv(v, x) + from scipy.special import iv + + return iv(v, x) def impl(self, v, x): return self.st_impl(v, x) @@ -1207,7 +1249,9 @@ class I1(UnaryScalarOp): @staticmethod def st_impl(x): - return scipy.special.i1(x) + from scipy.special import i1 + + return i1(x) def impl(self, x): return self.st_impl(x) @@ -1233,7 +1277,9 @@ class I0(UnaryScalarOp): @staticmethod def st_impl(x): - return scipy.special.i0(x) + from scipy.special import i0 + + return i0(x) def impl(self, x): return self.st_impl(x) @@ -1259,7 +1305,9 @@ class Ive(BinaryScalarOp): @staticmethod def st_impl(v, x): - return scipy.special.ive(v, x) + from scipy.special import ive + + return ive(v, x) def impl(self, v, x): return self.st_impl(v, x) @@ -1288,7 +1336,9 @@ class Kve(BinaryScalarOp): @staticmethod def st_impl(v, x): - return scipy.special.kve(v, x) + from scipy.special import kve + + return kve(v, x) def impl(self, v, x): return self.st_impl(v, x) @@ -1321,7 +1371,9 @@ class Sigmoid(UnaryScalarOp): nfunc_spec = ("scipy.special.expit", 1, 1) def impl(self, x): - return scipy.special.expit(x) + from scipy.special import expit + + return expit(x) def grad(self, inp, grads): (x,) = inp @@ -1496,7 +1548,9 @@ class BetaInc(ScalarOp): nfunc_spec = ("scipy.special.betainc", 3, 1) def impl(self, a, b, x): - return scipy.special.betainc(a, b, x) + from scipy.special import betainc + + return betainc(a, b, x) def grad(self, inp, grads): a, b, x = inp @@ -1756,7 +1810,9 @@ class BetaIncInv(ScalarOp): nfunc_spec = ("scipy.special.betaincinv", 3, 1) def impl(self, a, b, x): - return scipy.special.betaincinv(a, b, x) + from scipy.special import betaincinv + + return betaincinv(a, b, x) def grad(self, inputs, grads): (a, b, x) = inputs @@ -1796,7 +1852,9 @@ class Hyp2F1(ScalarOp): @staticmethod def st_impl(a, b, c, z): - return scipy.special.hyp2f1(a, b, c, z) + from scipy.special import hyp2f1 + + return hyp2f1(a, b, c, z) def impl(self, a, b, c, z): return Hyp2F1.st_impl(a, b, c, z) diff --git a/pytensor/scan/op.py b/pytensor/scan/op.py index bfe04a94d7..a01347ef9c 100644 --- a/pytensor/scan/op.py +++ b/pytensor/scan/op.py @@ -74,7 +74,6 @@ from pytensor.graph.replace import clone_replace from pytensor.graph.utils import InconsistencyError, MissingInputError from pytensor.link.c.basic import CLinker -from pytensor.link.c.exceptions import MissingGXX from pytensor.printing import op_debug_information from pytensor.scan.utils import ScanProfileStats, Validator, forced_replace, safe_new from pytensor.tensor.basic import as_tensor_variable @@ -1499,6 +1498,7 @@ def make_thunk(self, node, storage_map, compute_map, no_recycling, impl=None): then it must not do so for variables in the no_recycling list. """ + from pytensor.link.c.exceptions import MissingGXX # Before building the thunk, validate that the inner graph is # coherent diff --git a/pytensor/sparse/basic.py b/pytensor/sparse/basic.py index c590bc804a..4eeff11cce 100644 --- a/pytensor/sparse/basic.py +++ b/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 diff --git a/pytensor/sparse/rewriting.py b/pytensor/sparse/rewriting.py index bf6d6f0bc6..c32af1b3e6 100644 --- a/pytensor/sparse/rewriting.py +++ b/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) diff --git a/pytensor/sparse/sharedvar.py b/pytensor/sparse/sharedvar.py index 60b09656be..2ff7edf287 100644 --- a/pytensor/sparse/sharedvar.py +++ 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 ): diff --git a/pytensor/sparse/type.py b/pytensor/sparse/type.py index bbc8a9fda1..334da061d4 100644 --- a/pytensor/sparse/type.py +++ b/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) diff --git a/pytensor/tensor/blas.py b/pytensor/tensor/blas.py index 6170a02a98..3c38a9c501 100644 --- a/pytensor/tensor/blas.py +++ b/pytensor/tensor/blas.py @@ -111,50 +111,19 @@ _logger = logging.getLogger("pytensor.tensor.blas") -try: - import scipy.linalg.blas - - 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 - # when config.blas__ldflags is defined. So we don't need a - # warning in that case. - if not config.blas__ldflags: - _logger.warning( - "Failed to import scipy.linalg.blas, and " - "PyTensor flag blas__ldflags is empty. " - "Falling back on slower implementations for " - "dot(matrix, vector), dot(vector, matrix) and " - f"dot(vector, vector) ({e!s})" - ) - # If check_init_y() == True we need to initialize y when beta == 0. def check_init_y(): + # TODO: What is going on here? + from scipy.linalg.blas import get_blas_funcs + if check_init_y._result is None: - if not have_fblas: # pragma: no cover - check_init_y._result = False - else: - 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() + y = float("NaN") * np.ones((2,)) + x = np.ones((2,)) + A = np.ones((2, 2)) + gemv = get_blas_funcs("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 @@ -211,14 +180,15 @@ def make_node(self, y, alpha, A, x, beta): return Apply(self, inputs, [y.type()]) def perform(self, node, inputs, out_storage): + from scipy.linalg.blas import get_blas_funcs + y, alpha, A, x, beta = inputs if ( - have_fblas - and y.shape[0] != 0 + y.shape[0] != 0 and x.shape[0] != 0 - and y.dtype in _blas_gemv_fns + and y.dtype in {"float32", "float64", "complex64", "complex128"} ): - gemv = _blas_gemv_fns[y.dtype] + gemv = get_blas_funcs("gemv", dtype=y.dtype) if A.shape[0] != y.shape[0] or A.shape[1] != x.shape[0]: raise ValueError( diff --git a/pytensor/tensor/blas_scipy.py b/pytensor/tensor/blas_scipy.py index 16fb90988b..bb3ccf9354 100644 --- a/pytensor/tensor/blas_scipy.py +++ b/pytensor/tensor/blas_scipy.py @@ -2,30 +2,19 @@ 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 pytensor.tensor.blas import Ger class ScipyGer(Ger): def perform(self, node, inputs, output_storage): + from scipy.linalg.blas import get_blas_funcs + 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 = _blas_ger_fns[cA.dtype] + local_ger = get_blas_funcs("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 diff --git a/pytensor/tensor/random/basic.py b/pytensor/tensor/random/basic.py index bebcad55be..36d3842fd0 100644 --- a/pytensor/tensor/random/basic.py +++ b/pytensor/tensor/random/basic.py @@ -2,7 +2,6 @@ import warnings import numpy as np -import scipy.stats as stats import pytensor from pytensor.tensor import get_vector_length, specify_shape @@ -350,7 +349,9 @@ def rng_fn_scipy(cls, rng, loc, scale, size): is returned. """ - return stats.halfnorm.rvs(loc, scale, random_state=rng, size=size) + from scipy.stats import halfnorm + + return halfnorm.rvs(loc, scale, random_state=rng, size=size) halfnormal = HalfNormalRV() @@ -583,7 +584,9 @@ def __call__(self, b, scale=1.0, size=None, **kwargs): @classmethod def rng_fn_scipy(cls, rng, b, scale, size): - return stats.pareto.rvs(b, scale=scale, size=size, random_state=rng) + from scipy.stats import pareto + + return pareto.rvs(b, scale=scale, size=size, random_state=rng) pareto = ParetoRV() @@ -645,7 +648,9 @@ def rng_fn_scipy( scale: np.ndarray | float, size: list[int] | int | None, ) -> np.ndarray: - return stats.gumbel_r.rvs(loc=loc, scale=scale, size=size, random_state=rng) + from scipy.stats import gumbel_r + + return gumbel_r.rvs(loc=loc, scale=scale, size=size, random_state=rng) gumbel = GumbelRV() @@ -840,8 +845,10 @@ def safe_multivariate_normal(mean, cov, size=None, rng=None): variable. """ + from scipy.stats import multivariate_normal + res = np.atleast_1d( - stats.multivariate_normal(mean=mean, cov=cov, allow_singular=True).rvs( + multivariate_normal(mean=mean, cov=cov, allow_singular=True).rvs( size=size, random_state=rng ) ) @@ -1173,7 +1180,9 @@ def __call__(self, loc=0.0, scale=1.0, size=None, **kwargs): @classmethod def rng_fn_scipy(cls, rng, loc, scale, size): - return stats.cauchy.rvs(loc=loc, scale=scale, random_state=rng, size=size) + from scipy.stats import cauchy + + return cauchy.rvs(loc=loc, scale=scale, random_state=rng, size=size) cauchy = CauchyRV() @@ -1223,7 +1232,9 @@ def __call__(self, loc=0.0, scale=1.0, size=None, **kwargs): @classmethod def rng_fn_scipy(cls, rng, loc, scale, size): - return stats.halfcauchy.rvs(loc=loc, scale=scale, random_state=rng, size=size) + from scipy.stats import halfcauchy + + return halfcauchy.rvs(loc=loc, scale=scale, random_state=rng, size=size) halfcauchy = HalfCauchyRV() @@ -1277,7 +1288,9 @@ def __call__(self, shape, scale, size=None, **kwargs): @classmethod def rng_fn_scipy(cls, rng, shape, scale, size): - return stats.invgamma.rvs(shape, scale=scale, size=size, random_state=rng) + from scipy.stats import invgamma + + return invgamma.rvs(shape, scale=scale, size=size, random_state=rng) invgamma = InvGammaRV() @@ -1376,9 +1389,9 @@ def __call__(self, b, loc=0.0, scale=1.0, size=None, **kwargs): @classmethod def rng_fn_scipy(cls, rng, b, loc, scale, size): - return stats.truncexpon.rvs( - b, loc=loc, scale=scale, size=size, random_state=rng - ) + from scipy.stats import truncexpon + + return truncexpon.rvs(b, loc=loc, scale=scale, size=size, random_state=rng) truncexpon = TruncExponentialRV() @@ -1432,7 +1445,9 @@ def __call__(self, df, loc=0.0, scale=1.0, size=None, **kwargs): @classmethod def rng_fn_scipy(cls, rng, df, loc, scale, size): - return stats.t.rvs(df, loc=loc, scale=scale, size=size, random_state=rng) + from scipy.stats import t + + return t.rvs(df, loc=loc, scale=scale, size=size, random_state=rng) t = StudentTRV() @@ -1485,7 +1500,9 @@ def __call__(self, p, size=None, **kwargs): @classmethod def rng_fn_scipy(cls, rng, p, size): - return stats.bernoulli.rvs(p, size=size, random_state=rng) + from scipy.stats import bernoulli + + return bernoulli.rvs(p, size=size, random_state=rng) bernoulli = BernoulliRV() @@ -1624,7 +1641,9 @@ def __call__(self, n, p, size=None, **kwargs): @classmethod def rng_fn_scipy(cls, rng, n, p, size): - return stats.nbinom.rvs(n, p, size=size, random_state=rng) + from scipy.stats import nbinom + + return nbinom.rvs(n, p, size=size, random_state=rng) nbinom = NegBinomialRV() @@ -1681,7 +1700,9 @@ def __call__(self, n, a, b, size=None, **kwargs): @classmethod def rng_fn_scipy(cls, rng, n, a, b, size): - return stats.betabinom.rvs(n, a, b, size=size, random_state=rng) + from scipy.stats import betabinom + + return betabinom.rvs(n, a, b, size=size, random_state=rng) betabinom = BetaBinomialRV() @@ -1733,9 +1754,9 @@ def __call__(self, alpha=1.0, p=1.0, lambd=1.0, size=None, **kwargs): @classmethod def rng_fn_scipy(cls, rng, alpha, p, lambd, size): - return stats.gengamma.rvs( - alpha / p, p, scale=lambd, size=size, random_state=rng - ) + from scipy.stats import gengamma + + return gengamma.rvs(alpha / p, p, scale=lambd, size=size, random_state=rng) gengamma = GenGammaRV() diff --git a/pytensor/tensor/rewriting/blas_scipy.py b/pytensor/tensor/rewriting/blas_scipy.py index 610ef9b82f..2ed0279e45 100644 --- a/pytensor/tensor/rewriting/blas_scipy.py +++ b/pytensor/tensor/rewriting/blas_scipy.py @@ -1,5 +1,5 @@ from pytensor.graph.rewriting.basic import in2out -from pytensor.tensor.blas import ger, ger_destructive, have_fblas +from pytensor.tensor.blas import ger, ger_destructive from pytensor.tensor.blas_scipy import scipy_ger_inplace, scipy_ger_no_inplace from pytensor.tensor.rewriting.blas import blas_optdb, node_rewriter, optdb @@ -19,19 +19,19 @@ def make_ger_destructive(fgraph, node): use_scipy_blas = in2out(use_scipy_ger) make_scipy_blas_destructive = in2out(make_ger_destructive) -if have_fblas: - # scipy_blas is scheduled in the blas_optdb very late, because scipy sortof - # sucks, but it is almost always present. - # C implementations should be scheduled earlier than this, so that they take - # precedence. Once the original Ger is replaced, then these optimizations - # have no effect. - blas_optdb.register("scipy_blas", use_scipy_blas, "fast_run", position=100) - - # this matches the InplaceBlasOpt defined in blas.py - optdb.register( - "make_scipy_blas_destructive", - make_scipy_blas_destructive, - "fast_run", - "inplace", - position=50.2, - ) + +# scipy_blas is scheduled in the blas_optdb very late, because scipy sortof +# sucks [citation needed], but it is almost always present. +# C implementations should be scheduled earlier than this, so that they take +# precedence. Once the original Ger is replaced, then these optimizations +# have no effect. +blas_optdb.register("scipy_blas", use_scipy_blas, "fast_run", position=100) + +# this matches the InplaceBlasOpt defined in blas.py +optdb.register( + "make_scipy_blas_destructive", + make_scipy_blas_destructive, + "fast_run", + "inplace", + position=50.2, +) diff --git a/pytensor/tensor/slinalg.py b/pytensor/tensor/slinalg.py index 325567918a..accbc0a6b4 100644 --- a/pytensor/tensor/slinalg.py +++ b/pytensor/tensor/slinalg.py @@ -5,7 +5,6 @@ from typing import Literal, cast import numpy as np -import scipy.linalg import pytensor import pytensor.tensor as pt @@ -52,37 +51,41 @@ def infer_shape(self, fgraph, node, shapes): return [shapes[0]] def make_node(self, x): + from scipy.linalg import cholesky + x = as_tensor_variable(x) if x.type.ndim != 2: raise TypeError( f"Cholesky only allowed on matrix (2-D) inputs, got {x.type.ndim}-D input" ) # Call scipy to find output dtype - dtype = scipy.linalg.cholesky(np.eye(1, dtype=x.type.dtype)).dtype + dtype = cholesky(np.eye(1, dtype=x.type.dtype)).dtype return Apply(self, [x], [tensor(shape=x.type.shape, dtype=dtype)]) def perform(self, node, inputs, outputs): + from scipy.linalg import LinAlgError, cholesky + [x] = inputs [out] = outputs try: # Scipy cholesky only makes use of overwrite_a when it is F_CONTIGUOUS # If we have a `C_CONTIGUOUS` array we transpose to benefit from it if self.overwrite_a and x.flags["C_CONTIGUOUS"]: - out[0] = scipy.linalg.cholesky( + out[0] = cholesky( x.T, lower=not self.lower, check_finite=self.check_finite, overwrite_a=True, ).T else: - out[0] = scipy.linalg.cholesky( + out[0] = cholesky( x, lower=self.lower, check_finite=self.check_finite, overwrite_a=self.overwrite_a, ) - except scipy.linalg.LinAlgError: + except LinAlgError: if self.on_error == "raise": raise else: @@ -333,8 +336,10 @@ def __init__(self, **kwargs): super().__init__(**kwargs) def perform(self, node, inputs, output_storage): + from scipy.linalg import cho_solve + C, b = inputs - rval = scipy.linalg.cho_solve( + rval = cho_solve( (C, self.lower), b, check_finite=self.check_finite, @@ -400,8 +405,10 @@ def __init__(self, *, trans=0, unit_diagonal=False, **kwargs): self.unit_diagonal = unit_diagonal def perform(self, node, inputs, outputs): + from scipy.linalg import solve_triangular + A, b = inputs - outputs[0][0] = scipy.linalg.solve_triangular( + outputs[0][0] = solve_triangular( A, b, lower=self.lower, @@ -501,8 +508,10 @@ def __init__(self, *, assume_a="gen", **kwargs): self.assume_a = assume_a def perform(self, node, inputs, outputs): + from scipy.linalg import solve + a, b = inputs - outputs[0][0] = scipy.linalg.solve( + outputs[0][0] = solve( a=a, b=b, lower=self.lower, @@ -617,11 +626,13 @@ def make_node(self, a, b): return Apply(self, [a, b], [w]) def perform(self, node, inputs, outputs): + from scipy.linalg import eigvalsh + (w,) = outputs if len(inputs) == 2: - w[0] = scipy.linalg.eigvalsh(a=inputs[0], b=inputs[1], lower=self.lower) + w[0] = eigvalsh(a=inputs[0], b=inputs[1], lower=self.lower) else: - w[0] = scipy.linalg.eigvalsh(a=inputs[0], b=None, lower=self.lower) + w[0] = eigvalsh(a=inputs[0], b=None, lower=self.lower) def grad(self, inputs, g_outputs): a, b = inputs @@ -674,8 +685,10 @@ def make_node(self, a, b, gw): return Apply(self, [a, b, gw], [out1, out2]) def perform(self, node, inputs, outputs): + from scipy.linalg import eigh + (a, b, gw) = inputs - w, v = scipy.linalg.eigh(a, b, lower=self.lower) + w, v = eigh(a, b, lower=self.lower) gA = v.dot(np.diag(gw).dot(v.T)) gB = -v.dot(np.diag(gw * w).dot(v.T)) @@ -716,9 +729,11 @@ def make_node(self, A): ) def perform(self, node, inputs, outputs): + from scipy.linalg import expm + (A,) = inputs - (expm,) = outputs - expm[0] = scipy.linalg.expm(A) + (out,) = outputs + out[0] = expm(A) def grad(self, inputs, outputs): (A,) = inputs @@ -753,13 +768,15 @@ def infer_shape(self, fgraph, node, shapes): return [shapes[0]] def perform(self, node, inputs, outputs): + from scipy.linalg import eig, inv + # Kalbfleisch and Lawless, J. Am. Stat. Assoc. 80 (1985) Equation 3.4 # Kind of... You need to do some algebra from there to arrive at # this expression. (A, gA) = inputs (out,) = outputs - w, V = scipy.linalg.eig(A, right=True) - U = scipy.linalg.inv(V).T + w, V = eig(A, right=True) + U = inv(V).T exp_w = np.exp(w) X = np.subtract.outer(exp_w, exp_w) / np.subtract.outer(w, w) @@ -796,11 +813,13 @@ def make_node(self, A, B): return pytensor.graph.basic.Apply(self, [A, B], [X]) def perform(self, node, inputs, output_storage): + from scipy.linalg import solve_continuous_lyapunov + (A, B) = inputs X = output_storage[0] out_dtype = node.outputs[0].type.dtype - X[0] = scipy.linalg.solve_continuous_lyapunov(A, B).astype(out_dtype) + X[0] = solve_continuous_lyapunov(A, B).astype(out_dtype) def infer_shape(self, fgraph, node, shapes): return [shapes[0]] @@ -866,13 +885,13 @@ def make_node(self, A, B): return pytensor.graph.basic.Apply(self, [A, B], [X]) def perform(self, node, inputs, output_storage): + from scipy.linalg import solve_discrete_lyapunov + (A, B) = inputs X = output_storage[0] out_dtype = node.outputs[0].type.dtype - X[0] = scipy.linalg.solve_discrete_lyapunov(A, B, method="bilinear").astype( - out_dtype - ) + X[0] = solve_discrete_lyapunov(A, B, method="bilinear").astype(out_dtype) def infer_shape(self, fgraph, node, shapes): return [shapes[0]] @@ -985,6 +1004,8 @@ def make_node(self, A, B, Q, R): return pytensor.graph.basic.Apply(self, [A, B, Q, R], [X]) def perform(self, node, inputs, output_storage): + from scipy.linalg import solve_discrete_are + A, B, Q, R = inputs X = output_storage[0] @@ -992,7 +1013,7 @@ def perform(self, node, inputs, output_storage): Q = 0.5 * (Q + Q.T) out_dtype = node.outputs[0].type.dtype - X[0] = scipy.linalg.solve_discrete_are(A, B, Q, R).astype(out_dtype) + X[0] = solve_discrete_are(A, B, Q, R).astype(out_dtype) def infer_shape(self, fgraph, node, shapes): return [shapes[0]] @@ -1117,8 +1138,10 @@ def make_node(self, *matrices): return Apply(self, matrices, [out_type]) def perform(self, node, inputs, output_storage, params=None): + from scipy.linalg import block_diag + dtype = node.outputs[0].type.dtype - output_storage[0][0] = scipy.linalg.block_diag(*inputs).astype(dtype) + output_storage[0][0] = block_diag(*inputs).astype(dtype) def block_diag(*matrices: TensorVariable): diff --git a/pytensor/tensor/special.py b/pytensor/tensor/special.py index a2f02fabd8..2149ad5cb3 100644 --- a/pytensor/tensor/special.py +++ b/pytensor/tensor/special.py @@ -1,7 +1,6 @@ from textwrap import dedent import numpy as np -import scipy from pytensor.graph.basic import Apply from pytensor.graph.replace import _vectorize_node @@ -267,9 +266,11 @@ def make_node(self, x): return Apply(self, [x], [x.type()]) def perform(self, node, input_storage, output_storage): + from scipy.special import softmax + (x,) = input_storage (z,) = output_storage - z[0] = scipy.special.softmax(x, axis=self.axis) + z[0] = softmax(x, axis=self.axis) def L_op(self, inp, outputs, grads): (x,) = inp @@ -519,9 +520,11 @@ def make_node(self, x): return Apply(self, [x], [x.type()]) def perform(self, node, input_storage, output_storage): + from scipy.special import log_softmax + (x,) = input_storage (z,) = output_storage - z[0] = scipy.special.log_softmax(x, axis=self.axis) + z[0] = log_softmax(x, axis=self.axis) def grad(self, inp, grads): (x,) = inp diff --git a/tests/d3viz/test_d3viz.py b/tests/d3viz/test_d3viz.py index b6b6479a1b..7e4b0426a0 100644 --- a/tests/d3viz/test_d3viz.py +++ b/tests/d3viz/test_d3viz.py @@ -9,12 +9,14 @@ from pytensor import compile from pytensor.compile.function import function from pytensor.configdefaults import config -from pytensor.printing import pydot_imported, pydot_imported_msg +from pytensor.printing import _try_pydot_import from tests.d3viz import models -if not pydot_imported: - pytest.skip("pydot not available: " + pydot_imported_msg, allow_module_level=True) +try: + _try_pydot_import() +except Exception as e: + pytest.skip(f"pydot not available: {e!s}", allow_module_level=True) class TestD3Viz: diff --git a/tests/d3viz/test_formatting.py b/tests/d3viz/test_formatting.py index 9f5f8be9ec..7d1149be0e 100644 --- a/tests/d3viz/test_formatting.py +++ b/tests/d3viz/test_formatting.py @@ -3,11 +3,13 @@ from pytensor import config, function from pytensor.d3viz.formatting import PyDotFormatter -from pytensor.printing import pydot_imported, pydot_imported_msg +from pytensor.printing import _try_pydot_import -if not pydot_imported: - pytest.skip("pydot not available: " + pydot_imported_msg, allow_module_level=True) +try: + _try_pydot_import() +except Exception as e: + pytest.skip(f"pydot not available: {e!s}", allow_module_level=True) from tests.d3viz import models diff --git a/tests/scan/test_printing.py b/tests/scan/test_printing.py index 44465f0152..9bf32af48f 100644 --- a/tests/scan/test_printing.py +++ b/tests/scan/test_printing.py @@ -5,7 +5,7 @@ import pytensor.tensor as pt from pytensor.configdefaults import config from pytensor.graph.fg import FunctionGraph -from pytensor.printing import debugprint, pydot_imported, pydotprint +from pytensor.printing import _try_pydot_import, debugprint, pydotprint from pytensor.tensor.type import dvector, iscalar, scalar, vector @@ -686,6 +686,13 @@ def no_shared_fn(n, x_tm1, M): assert truth.strip() == out.strip() +try: + _try_pydot_import() + pydot_imported = True +except Exception: + pydot_imported = False + + @pytest.mark.skipif(not pydot_imported, reason="pydot not available") def test_pydotprint(): def f_pow2(x_tm1): diff --git a/tests/sparse/test_basic.py b/tests/sparse/test_basic.py index 4075ed3ed6..986c95230d 100644 --- a/tests/sparse/test_basic.py +++ b/tests/sparse/test_basic.py @@ -3,6 +3,7 @@ import numpy as np import pytest +import scipy as sp from packaging import version import pytensor @@ -82,7 +83,6 @@ _is_dense_variable, _is_sparse, _is_sparse_variable, - _mtypes, ) from pytensor.sparse.rewriting import ( AddSD_ccode, @@ -115,8 +115,8 @@ from tests.tensor.test_sharedvar import makeSharedTester -sp = pytest.importorskip("scipy", minversion="0.7.0") - +sparse_formats = ["csc", "csr"] +_mtypes = [sp.sparse.csc_matrix, sp.sparse.csr_matrix] # Probability distributions are currently tested in test_sp2.py # from pytensor.sparse import ( @@ -1997,7 +1997,7 @@ def setup_method(self): self.op = sparse.col_scale def test_op(self): - for format in sparse.sparse_formats: + for format in sparse_formats: variable, data = sparse_random_inputs(format, shape=(8, 10)) variable.append(vector()) data.append(np.random.random(10).astype(config.floatX)) @@ -2020,7 +2020,7 @@ def test_infer_shape(self): self._compile_and_check(variable, [self.op(*variable)], data, cls) def test_grad(self): - for format in sparse.sparse_formats: + for format in sparse_formats: variable, data = sparse_random_inputs(format, shape=(8, 10)) variable.append(vector()) data.append(np.random.random(10).astype(config.floatX)) @@ -2034,7 +2034,7 @@ def setup_method(self): self.op = sparse.row_scale def test_op(self): - for format in sparse.sparse_formats: + for format in sparse_formats: variable, data = sparse_random_inputs(format, shape=(8, 10)) variable.append(vector()) data.append(np.random.random(8).astype(config.floatX)) @@ -2057,7 +2057,7 @@ def test_infer_shape(self): self._compile_and_check(variable, [self.op(*variable)], data, cls) def test_grad(self): - for format in sparse.sparse_formats: + for format in sparse_formats: variable, data = sparse_random_inputs(format, shape=(8, 10)) variable.append(vector()) data.append(np.random.random(8).astype(config.floatX)) @@ -2075,7 +2075,7 @@ def setup_method(self): @pytest.mark.parametrize("op_type", ["func", "method"]) def test_op(self, op_type): - for format in sparse.sparse_formats: + for format in sparse_formats: for axis in self.possible_axis: variable, data = sparse_random_inputs(format, shape=(10, 10)) @@ -2095,7 +2095,7 @@ def test_op(self, op_type): utt.assert_allclose(expected, tested) def test_infer_shape(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for axis in self.possible_axis: variable, data = sparse_random_inputs(format, shape=(9, 10)) self._compile_and_check( @@ -2103,7 +2103,7 @@ def test_infer_shape(self): ) def test_grad(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for axis in self.possible_axis: for struct in [True, False]: variable, data = sparse_random_inputs(format, shape=(9, 10)) @@ -2121,7 +2121,7 @@ def setup_method(self): self.op = diag def test_op(self): - for format in sparse.sparse_formats: + for format in sparse_formats: variable, data = sparse_random_inputs(format, shape=(10, 10)) z = self.op(*variable) @@ -2134,14 +2134,14 @@ def test_op(self): utt.assert_allclose(expected, tested) def test_infer_shape(self): - for format in sparse.sparse_formats: + for format in sparse_formats: variable, data = sparse_random_inputs(format, shape=(10, 10)) self._compile_and_check( variable, [self.op(*variable)], data, self.op_class, warn=False ) def test_grad(self): - for format in sparse.sparse_formats: + for format in sparse_formats: variable, data = sparse_random_inputs(format, shape=(10, 10)) verify_grad_sparse(self.op, data, structured=False) @@ -2153,7 +2153,7 @@ def setup_method(self): self.op = square_diagonal def test_op(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for size in range(5, 9): variable = [vector()] data = [np.random.random(size).astype(config.floatX)] @@ -2167,7 +2167,7 @@ def test_op(self): assert tested.shape == expected.shape def test_infer_shape(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for size in range(5, 9): variable = [vector()] data = [np.random.random(size).astype(config.floatX)] @@ -2177,7 +2177,7 @@ def test_infer_shape(self): ) def test_grad(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for size in range(5, 9): data = [np.random.random(size).astype(config.floatX)] @@ -2191,7 +2191,7 @@ def setup_method(self): self.op = ensure_sorted_indices def test_op(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for shape in zip(range(5, 9), range(3, 7)[::-1], strict=True): variable, data = sparse_random_inputs(format, shape=shape) @@ -2202,7 +2202,7 @@ def test_op(self): utt.assert_allclose(expected, tested) def test_infer_shape(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for shape in zip(range(5, 9), range(3, 7)[::-1], strict=True): variable, data = sparse_random_inputs(format, shape=shape) self._compile_and_check( @@ -2210,7 +2210,7 @@ def test_infer_shape(self): ) def test_grad(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for shape in zip(range(5, 9), range(3, 7)[::-1], strict=True): variable, data = sparse_random_inputs(format, shape=shape) verify_grad_sparse(self.op, data, structured=False) @@ -2222,7 +2222,7 @@ def setup_method(self): self.op = clean def test_op(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for shape in zip(range(5, 9), range(3, 7)[::-1], strict=True): variable, data = sparse_random_inputs(format, shape=shape) @@ -2241,7 +2241,7 @@ def test_op(self): utt.assert_allclose(expected, tested) def test_grad(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for shape in zip(range(5, 9), range(3, 7)[::-1], strict=True): variable, data = sparse_random_inputs(format, shape=shape) verify_grad_sparse(self.op, data, structured=False) @@ -2629,7 +2629,7 @@ def setup_method(self): # slow but only test def test_cast(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for i_dtype in sparse.all_dtypes: for o_dtype in sparse.all_dtypes: (variable,), (data,) = sparse_random_inputs( @@ -2658,7 +2658,7 @@ def test_cast(self): @pytest.mark.slow def test_infer_shape(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for i_dtype in sparse.all_dtypes: for o_dtype in sparse.all_dtypes: variable, data = sparse_random_inputs( @@ -2669,7 +2669,7 @@ def test_infer_shape(self): ) def test_grad(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for i_dtype in sparse.float_dtypes: for o_dtype in float_dtypes: if o_dtype == "float16": @@ -2690,7 +2690,7 @@ def _format_info(nb): x = {} mat = {} - for format in sparse.sparse_formats: + for format in sparse_formats: variable = getattr(pytensor.sparse, format + "_matrix") spa = getattr(sp.sparse, format + "_matrix") @@ -2710,8 +2710,8 @@ class _TestHVStack(utt.InferShapeTester): x, mat = _format_info(nb) def test_op(self): - for format in sparse.sparse_formats: - for out_f in sparse.sparse_formats: + for format in sparse_formats: + for out_f in sparse_formats: for dtype in sparse.all_dtypes: blocks = self.mat[format] @@ -2729,7 +2729,7 @@ def test_op(self): assert tested.dtype == expected.dtype def test_infer_shape(self): - for format in sparse.sparse_formats: + for format in sparse_formats: self._compile_and_check( self.x[format], [self.op_class(dtype="float64")(*self.x[format])], @@ -2738,8 +2738,8 @@ def test_infer_shape(self): ) def test_grad(self): - for format in sparse.sparse_formats: - for out_f in sparse.sparse_formats: + for format in sparse_formats: + for out_f in sparse_formats: for dtype in sparse.float_dtypes: verify_grad_sparse( self.op_class(format=out_f, dtype=dtype), @@ -2782,7 +2782,7 @@ def setup_method(self): super().setup_method() self.op_class = AddSSData - for format in sparse.sparse_formats: + for format in sparse_formats: variable = getattr(pytensor.sparse, format + "_matrix") a_val = np.array( @@ -2795,7 +2795,7 @@ def setup_method(self): self.a[format] = [constant for t in range(2)] def test_op(self): - for format in sparse.sparse_formats: + for format in sparse_formats: f = pytensor.function(self.x[format], add_s_s_data(*self.x[format])) tested = f(*self.a[format]) @@ -2806,7 +2806,7 @@ def test_op(self): assert tested.dtype == expected.dtype def test_infer_shape(self): - for format in sparse.sparse_formats: + for format in sparse_formats: self._compile_and_check( self.x[format], [add_s_s_data(*self.x[format])], @@ -2815,7 +2815,7 @@ def test_infer_shape(self): ) def test_grad(self): - for format in sparse.sparse_formats: + for format in sparse_formats: verify_grad_sparse(self.op_class(), self.a[format], structured=True) @@ -2864,7 +2864,7 @@ def setup_method(self): assert eval(self.__class__.__name__) is self.__class__ def test_op(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for dtype in test_dtypes: if dtype == "int8" or dtype == "uint8": continue @@ -2967,7 +2967,7 @@ def test_op(self): if grad_test: def test_grad(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for dtype in sparse.float_dtypes: variable, data = sparse_random_inputs( format, shape=(4, 7), out_dtype=dtype, gap=self.gap_grad @@ -3244,7 +3244,7 @@ def setup_method(self): self.op_class = TrueDot def test_op_ss(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for dtype in sparse.all_dtypes: variable, data = sparse_random_inputs( format, shape=(10, 10), out_dtype=dtype, n=2, p=0.1 @@ -3263,7 +3263,7 @@ def test_op_ss(self): utt.assert_allclose(tested, expected) def test_op_sd(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for dtype in sparse.all_dtypes: variable, data = sparse_random_inputs( format, shape=(10, 10), out_dtype=dtype, n=2, p=0.1 @@ -3282,7 +3282,7 @@ def test_op_sd(self): utt.assert_allclose(tested, expected) def test_infer_shape(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for dtype in sparse.all_dtypes: (x,), (x_value,) = sparse_random_inputs( format, shape=(9, 10), out_dtype=dtype, p=0.1 @@ -3297,7 +3297,7 @@ def test_infer_shape(self): ) def test_grad(self): - for format in sparse.sparse_formats: + for format in sparse_formats: for dtype in sparse.float_dtypes: (x,), (x_value,) = sparse_random_inputs( format, shape=(9, 10), out_dtype=dtype, p=0.1 diff --git a/tests/sparse/test_rewriting.py b/tests/sparse/test_rewriting.py index 280d9dbf70..009fd8e31a 100644 --- a/tests/sparse/test_rewriting.py +++ b/tests/sparse/test_rewriting.py @@ -14,6 +14,9 @@ from tests.sparse.test_basic import random_lil +sparse_formats = ["csc", "csr"] + + def test_local_csm_properties_csm(): data = vector() indices, indptr, shape = (ivector(), ivector(), ivector()) @@ -71,7 +74,7 @@ def test_local_mul_s_d(): mode = get_default_mode() mode = mode.including("specialize", "local_mul_s_d") - for sp_format in sparse.sparse_formats: + for sp_format in sparse_formats: inputs = [getattr(pytensor.sparse, sp_format + "_matrix")(), matrix()] f = pytensor.function(inputs, sparse.mul_s_d(*inputs), mode=mode) diff --git a/tests/sparse/test_sp2.py b/tests/sparse/test_sp2.py index 54b88f4491..e202d2890b 100644 --- a/tests/sparse/test_sp2.py +++ b/tests/sparse/test_sp2.py @@ -1,9 +1,5 @@ -import pytest - - -sp = pytest.importorskip("scipy", minversion="0.7.0") - import numpy as np +import scipy as sp import pytensor from pytensor import sparse @@ -20,11 +16,14 @@ from tests.sparse.test_basic import as_sparse_format +sparse_formats = ["csr", "csc"] + + class TestPoisson(utt.InferShapeTester): x = {} a = {} - for format in sparse.sparse_formats: + for format in sparse_formats: variable = getattr(pytensor.sparse, format + "_matrix") a_val = np.array( @@ -40,7 +39,7 @@ def setup_method(self): self.op_class = Poisson def test_op(self): - for format in sparse.sparse_formats: + for format in sparse_formats: f = pytensor.function([self.x[format]], poisson(self.x[format])) tested = f(self.a[format]) @@ -51,7 +50,7 @@ def test_op(self): assert tested.shape == self.a[format].shape def test_infer_shape(self): - for format in sparse.sparse_formats: + for format in sparse_formats: self._compile_and_check( [self.x[format]], [poisson(self.x[format])], @@ -76,7 +75,7 @@ def setup_method(self): self.op_class = Binomial def test_op(self): - for sp_format in sparse.sparse_formats: + for sp_format in sparse_formats: for o_type in sparse.float_dtypes: f = pytensor.function( self.inputs, Binomial(sp_format, o_type)(*self.inputs) @@ -90,7 +89,7 @@ def test_op(self): assert np.allclose(np.floor(tested.todense()), tested.todense()) def test_infer_shape(self): - for sp_format in sparse.sparse_formats: + for sp_format in sparse_formats: for o_type in sparse.float_dtypes: self._compile_and_check( self.inputs, diff --git a/tests/tensor/test_blas_scipy.py b/tests/tensor/test_blas_scipy.py index 7cdfaadc34..716eab7bbe 100644 --- a/tests/tensor/test_blas_scipy.py +++ b/tests/tensor/test_blas_scipy.py @@ -1,7 +1,6 @@ import pickle import numpy as np -import pytest import pytensor from pytensor import tensor as pt @@ -12,7 +11,6 @@ from tests.unittest_tools import OptimizationTestMixin -@pytest.mark.skipif(not pytensor.tensor.blas_scipy.have_fblas, reason="fblas needed") class TestScipyGer(OptimizationTestMixin): def setup_method(self): self.mode = pytensor.compile.get_default_mode() diff --git a/tests/test_printing.py b/tests/test_printing.py index 73403880e9..be5dbbc5a1 100644 --- a/tests/test_printing.py +++ b/tests/test_printing.py @@ -17,13 +17,13 @@ PatternPrinter, PPrinter, Print, + _try_pydot_import, char_from_number, debugprint, default_printer, get_node_by_id, min_informative_str, pp, - pydot_imported, pydotprint, ) from pytensor.tensor import as_tensor_variable @@ -31,6 +31,13 @@ from tests.graph.utils import MyInnerGraphOp, MyOp, MyVariable +try: + _try_pydot_import() + pydot_imported = True +except Exception: + pydot_imported = False + + @pytest.mark.parametrize( "number,s", [