diff --git a/pytensor/__init__.py b/pytensor/__init__.py index 12f67c9a37..77180abab9 100644 --- a/pytensor/__init__.py +++ b/pytensor/__init__.py @@ -1,168 +1,32 @@ -""" -PyTensor is an optimizing compiler in Python, built to evaluate -complicated expressions (especially matrix-valued ones) as quickly as -possible. PyTensor compiles expression graphs (see :doc:`graph` ) that -are built by Python code. The expressions in these graphs are called -`Apply` nodes and the variables in these graphs are called `Variable` -nodes. - -You compile a graph by calling `function`, which takes a graph, and -returns a callable object. One of pytensor's most important features is -that `function` can transform your graph before compiling it. It can -replace simple expressions with faster or more numerically stable -implementations. - -To learn more, check out: - -- Op List (:doc:`oplist`) - -""" - __docformat__ = "restructuredtext en" -# Set a default logger. It is important to do this before importing some other -# pytensor code, since this code may want to log some messages. -import logging -import sys -import warnings -from functools import singledispatch -from pathlib import Path -from typing import Any, NoReturn, Optional from pytensor import _version __version__: str = _version.get_versions()["version"] +del _version -pytensor_logger = logging.getLogger("pytensor") -logging_default_handler = logging.StreamHandler() -logging_default_formatter = logging.Formatter( - fmt="%(levelname)s (%(name)s): %(message)s" -) -logging_default_handler.setFormatter(logging_default_formatter) -pytensor_logger.setLevel(logging.WARNING) - -if not pytensor_logger.hasHandlers(): - pytensor_logger.addHandler(logging_default_handler) - - -# Disable default log handler added to pytensor_logger when the module -# is imported. -def disable_log_handler(logger=pytensor_logger, handler=logging_default_handler): - if logger.hasHandlers(): - logger.removeHandler(handler) - - -# Raise a meaningful warning/error if the pytensor directory is in the Python -# path. -rpath = Path(__file__).parent.resolve() -if any(rpath == Path(p).resolve() for p in sys.path): - raise RuntimeError("You have the pytensor directory in your Python path.") from pytensor.configdefaults import config -# This is the api version for ops that generate C code. External ops -# might need manual changes if this number goes up. An undefined -# __api_version__ can be understood to mean api version 0. -# -# This number is not tied to the release version and should change -# very rarely. -__api_version__ = 1 - -# isort: off -from pytensor.graph.basic import Variable -from pytensor.graph.replace import clone_replace, graph_replace - -# isort: on - - -def as_symbolic(x: Any, name: str | None = None, **kwargs) -> Variable: - """Convert `x` into an equivalent PyTensor `Variable`. - - Parameters - ---------- - x - The object to be converted into a ``Variable`` type. A - ``numpy.ndarray`` argument will not be copied, but a list of numbers - will be copied to make an ``numpy.ndarray``. - name - If a new ``Variable`` instance is created, it will be named with this - string. - kwargs - Options passed to the appropriate sub-dispatch functions. For example, - `ndim` and `dtype` can be passed when `x` is an `numpy.ndarray` or - `Number` type. - - Raises - ------ - TypeError - If `x` cannot be converted to a `Variable`. - - """ - if isinstance(x, Variable): - return x - - res = _as_symbolic(x, **kwargs) - res.name = name - return res - - -@singledispatch -def _as_symbolic(x: Any, **kwargs) -> Variable: - from pytensor.tensor import as_tensor_variable - - return as_tensor_variable(x, **kwargs) - - # isort: off -from pytensor import scalar, tensor +from pytensor import tensor +from pytensor import sparse from pytensor.compile import ( In, Mode, Out, - ProfileStats, - predefined_linkers, - predefined_modes, - predefined_optimizers, shared, + wrap_py, + function, ) -from pytensor.compile.function import function, function_dump -from pytensor.compile.function.types import FunctionMaker -from pytensor.gradient import Lop, Rop, grad, subgraph_grad +from pytensor.gradient import Lop, Rop, grad from pytensor.printing import debugprint as dprint -from pytensor.printing import pp, pprint -from pytensor.updates import OrderedUpdates - -# isort: on - - -def get_underlying_scalar_constant(v): - """Return the constant scalar (i.e. 0-D) value underlying variable `v`. - - If `v` is the output of dim-shuffles, fills, allocs, cast, etc. - this function digs through them. - - If ``pytensor.sparse`` is also there, we will look over CSM `Op`. - If `v` is not some view of constant data, then raise a - `NotScalarConstantError`. - """ - warnings.warn( - "get_underlying_scalar_constant is deprecated. Use tensor.get_underlying_scalar_constant_value instead.", - FutureWarning, - ) - from pytensor.tensor.basic import get_underlying_scalar_constant_value - - return get_underlying_scalar_constant_value(v) - - -# isort: off -import pytensor.tensor.random.var -import pytensor.sparse from pytensor.ifelse import ifelse -from pytensor.scan import checkpoints from pytensor.scan.basic import scan from pytensor.scan.views import foldl, foldr, map, reduce from pytensor.compile.builders import OpFromGraph diff --git a/pytensor/basic.py b/pytensor/basic.py new file mode 100644 index 0000000000..6b540f20d8 --- /dev/null +++ b/pytensor/basic.py @@ -0,0 +1,42 @@ +from functools import singledispatch +from typing import Any + +from pytensor.graph import Variable + + +def as_symbolic(x: Any, name: str | None = None, **kwargs) -> Variable: + """Convert `x` into an equivalent PyTensor `Variable`. + + Parameters + ---------- + x + The object to be converted into a ``Variable`` type. A + ``numpy.ndarray`` argument will not be copied, but a list of numbers + will be copied to make an ``numpy.ndarray``. + name + If a new ``Variable`` instance is created, it will be named with this + string. + kwargs + Options passed to the appropriate sub-dispatch functions. For example, + `ndim` and `dtype` can be passed when `x` is an `numpy.ndarray` or + `Number` type. + + Raises + ------ + TypeError + If `x` cannot be converted to a `Variable`. + + """ + if isinstance(x, Variable): + return x + + res = _as_symbolic(x, **kwargs) + res.name = name + return res + + +@singledispatch +def _as_symbolic(x: Any, **kwargs) -> Variable: + from pytensor.tensor import as_tensor_variable + + return as_tensor_variable(x, **kwargs) diff --git a/pytensor/compile/mode.py b/pytensor/compile/mode.py index 5a5e0c9cdc..ca7d5387fc 100644 --- a/pytensor/compile/mode.py +++ b/pytensor/compile/mode.py @@ -5,7 +5,7 @@ import logging import warnings -from typing import Any, Literal +from typing import Any from pytensor.compile.function.types import Supervisor from pytensor.configdefaults import config @@ -512,7 +512,7 @@ def get_mode(orig_string): if upper_string == "FAST_RUN": linker = config.linker if linker == "auto": - return CVM if config.cxx else VM + return NUMBA return fast_run_linkers_to_mode[linker] global _CACHED_RUNTIME_MODES @@ -565,26 +565,3 @@ def register_mode(name, mode): if name in predefined_modes: raise ValueError(f"Mode name already taken: {name}") predefined_modes[name] = mode - - -def get_target_language(mode=None) -> tuple[Literal["py", "c", "numba", "jax"], ...]: - """Get the compilation target language.""" - - if mode is None: - mode = get_default_mode() - - linker = mode.linker - - if isinstance(linker, NumbaLinker): - return ("numba",) - if isinstance(linker, JAXLinker): - return ("jax",) - if isinstance(linker, PerformLinker): - return ("py",) - if isinstance(linker, CLinker): - return ("c",) - - if isinstance(linker, VMLinker | OpWiseCLinker): - return ("c", "py") if config.cxx else ("py",) - - raise Exception(f"Unsupported Linker: {linker}") diff --git a/pytensor/configdefaults.py b/pytensor/configdefaults.py index 684ea94f7a..0689fbbdf8 100644 --- a/pytensor/configdefaults.py +++ b/pytensor/configdefaults.py @@ -10,11 +10,9 @@ import numpy as np -import pytensor from pytensor.configparser import ( BoolParam, ConfigParam, - DeviceParam, EnumStr, FloatParam, IntParam, @@ -67,15 +65,6 @@ def _filter_mode(val): ) -def _warn_cxx(val): - """We only support clang++ as otherwise we hit strange g++/OSX bugs.""" - if sys.platform == "darwin" and val and "clang++" not in val: - _logger.warning( - "Only clang++ is supported. With g++, we end up with strange g++/OSX bugs." - ) - return True - - def _split_version(version): """ Take version as a dot-separated string, return a tuple of int @@ -96,16 +85,6 @@ def _warn_default(version): return True -def _good_seem_param(seed): - if seed == "random": - return True - try: - int(seed) - except Exception: - return False - return True - - def _is_valid_check_preallocated_output_param(param): if not isinstance(param, str): return False @@ -263,13 +242,6 @@ def add_basic_configvars(): in_c_key=False, ) - config.add( - "device", - ("Default device for computations. only cpu is supported for now"), - DeviceParam("cpu", mutable=False), - in_c_key=False, - ) - config.add( "print_global_stats", "Print some global statistics (time spent) at the end", @@ -293,10 +265,6 @@ def _is_gt_0(x): return x > 0 -def _is_greater_or_equal_0(x): - return x >= 0 - - def add_compile_configvars(): config.add( "mode", @@ -305,91 +273,20 @@ def add_compile_configvars(): in_c_key=False, ) - param = "g++" - - # Test whether or not g++ is present: disable C code if it is not. - try: - rc = call_subprocess_Popen(["g++", "-v"]) - except OSError: - rc = 1 - - # Anaconda on Windows has mingw-w64 packages including GCC, but it may not be on PATH. - if rc != 0: - if sys.platform == "win32": - mingw_w64_gcc = Path(sys.executable).parent / "Library/mingw-w64/bin/g++" - try: - rc = call_subprocess_Popen([str(mingw_w64_gcc), "-v"]) - if rc == 0: - maybe_add_to_os_environ_pathlist("PATH", mingw_w64_gcc.parent) - except OSError: - rc = 1 - if rc != 0: - _logger.warning( - "g++ not available, if using conda: `conda install gxx`" - ) - - if rc != 0: - param = "" - - # On Mac/FreeBSD we test for 'clang++' and use it by default - if sys.platform == "darwin" or sys.platform.startswith("freebsd"): - try: - rc = call_subprocess_Popen(["clang++", "-v"]) - if rc == 0: - param = "clang++" - except OSError: - pass - - # Try to find the full compiler path from the name - if param != "": - newp = which(param) - if newp is not None: - param = newp - del newp - - # to support path that includes spaces, we need to wrap it with double quotes on Windows - if param and os.name == "nt": - param = f'"{param}"' - - config.add( - "cxx", - "The C++ compiler to use. Currently only g++ is" - " supported, but supporting additional compilers should not be " - "too difficult. " - "If it is empty, no C++ code is compiled.", - StrParam(param, validate=_warn_cxx), - in_c_key=False, - ) - del param - default_linker = "auto" - if rc == 0 and config.cxx != "": - # Keep the default linker the same as the one for the mode FAST_RUN - linker_options = [ - "cvm", - "c|py", - "py", - "c", - "c|py_nogc", - "vm", - "vm_nogc", - "cvm_nogc", - "numba", - "jax", - ] - else: - # g++ is not present or the user disabled it, - # linker should default to python only. - linker_options = ["py", "vm", "vm_nogc", "numba", "jax"] - if type(config).cxx.is_default: - # If the user provided an empty value for cxx, do not warn. - _logger.warning( - "g++ not detected! PyTensor will be unable to compile " - "C-implementations and will default to Python. " - "Performance may be severely degraded. " - "To remove this warning, set PyTensor flags cxx to an empty string." - ) + linker_options = [ + "cvm", + "c|py", + "py", + "c", + "c|py_nogc", + "vm", + "vm_nogc", + "cvm_nogc", + "numba", + "jax", + ] config.add( "linker", @@ -453,13 +350,6 @@ def add_compile_configvars(): in_c_key=False, ) - config.add( - "nocleanup", - "Suppress the deletion of code files that did not compile cleanly", - BoolParam(False), - in_c_key=False, - ) - config.add( "on_unused_input", "What to do if a variable in the 'inputs' list of " @@ -468,6 +358,96 @@ def add_compile_configvars(): in_c_key=False, ) + +def add_c_compile_configvars(): + def get_cxx_default(): + param = "g++" + + # Test whether or not g++ is present: disable C code if it is not. + try: + rc = call_subprocess_Popen(["g++", "-v"]) + except OSError: + rc = 1 + + # Anaconda on Windows has mingw-w64 packages including GCC, but it may not be on PATH. + if rc != 0: + if sys.platform == "win32": + mingw_w64_gcc = ( + Path(sys.executable).parent / "Library/mingw-w64/bin/g++" + ) + try: + rc = call_subprocess_Popen([str(mingw_w64_gcc), "-v"]) + if rc == 0: + maybe_add_to_os_environ_pathlist("PATH", mingw_w64_gcc.parent) + except OSError: + rc = 1 + if rc != 0: + _logger.warning( + "g++ not available, if using conda: `conda install gxx`" + ) + + if rc != 0: + param = "" + + # On Mac/FreeBSD we test for 'clang++' and use it by default + if sys.platform == "darwin" or sys.platform.startswith("freebsd"): + try: + rc = call_subprocess_Popen(["clang++", "-v"]) + if rc == 0: + param = "clang++" + except OSError: + pass + + # Try to find the full compiler path from the name + if param != "": + newp = which(param) + if newp is not None: + param = newp + del newp + + # to support path that includes spaces, we need to wrap it with double quotes on Windows + if param and os.name == "nt": + param = f'"{param}"' + + return param + + def warn_cxx(val): + """We only support clang++ as otherwise we hit strange g++/OSX bugs.""" + if sys.platform == "darwin" and val and "clang++" not in val: + _logger.warning( + "Only clang++ is supported. With g++, we end up with strange g++/OSX bugs." + ) + return True + + config.add( + "cxx", + "The C++ compiler to use. Currently only g++ is" + " supported, but supporting additional compilers should not be " + "too difficult. " + "If it is empty, no C++ code is compiled.", + StrParam(get_cxx_default, validate=warn_cxx), + in_c_key=False, + ) + + def default_gcc_version_str(): + if config.cxx != "": + try: + p_out = output_subprocess_Popen([config.cxx, "-dumpversion"]) + gcc_version_str = p_out[0].strip().decode() + except OSError: + # Typically means gcc cannot be found. + gcc_version_str = "GCC_NOT_FOUND" + else: + gcc_version_str = "GCC_NOT_FOUND" + return gcc_version_str + + config.add( + "gcc_version_str", + "", + StrParam(default_gcc_version_str, mutable=False), + in_c_key=False, + ) + config.add( "gcc__cxxflags", "Extra compiler flags for gcc", @@ -476,6 +456,13 @@ def add_compile_configvars(): in_c_key=False, ) + config.add( + "nocleanup", + "Suppress the deletion of code files that did not compile cleanly", + BoolParam(False), + in_c_key=False, + ) + config.add( "cmodule__warn_no_version", "If True, will print a warning when compiling one or more Op " @@ -540,15 +527,11 @@ def add_compile_configvars(): lock is held by the same owner *and* has not been 'refreshed' by this owner for more than this period. Refreshes are done every half timeout period for running processes.""", - IntParam(_timeout_default, validate=_is_greater_or_equal_0, mutable=False), + IntParam(_timeout_default, validate=lambda x: x >= 0, mutable=False), in_c_key=False, ) -def _is_valid_cmp_sloppy(v): - return v in (0, 1, 2) - - def add_tensor_configvars(): # This flag is used when we import PyTensor to initialize global variables. # So changing it after import will not modify these global variables. @@ -557,7 +540,7 @@ def add_tensor_configvars(): config.add( "tensor__cmp_sloppy", "Relax pytensor.tensor.math._allclose (0) not at all, (1) a bit, (2) more", - IntParam(0, _is_valid_cmp_sloppy, mutable=False), + IntParam(0, validate=lambda x: x in (0, 1, 2), mutable=False), in_c_key=False, ) @@ -664,14 +647,6 @@ def add_error_and_warning_configvars(): ) -def _has_cxx(): - return bool(config.cxx) - - -def _is_valid_check_strides(v): - return v in (0, 1, 2) - - def add_testvalue_and_checking_configvars(): config.add( "print_test_value", @@ -754,7 +729,7 @@ def add_testvalue_and_checking_configvars(): config.add( "DebugMode__check_c", "Run C implementations where possible", - BoolParam(_has_cxx), + BoolParam(lambda: bool(config.cxx)), in_c_key=False, ) @@ -779,7 +754,7 @@ def add_testvalue_and_checking_configvars(): "On difference: (0) - ignore, (1) warn, or (2) raise error" ), # TODO: make this an Enum setting - IntParam(0, _is_valid_check_strides), + IntParam(0, validate=lambda x: x in (0, 1, 2)), in_c_key=False, ) @@ -855,7 +830,7 @@ def add_testvalue_and_checking_configvars(): "profiling__min_memory_size", """For the memory profile, do not print Apply nodes if the size of their outputs (in bytes) is lower than this threshold""", - IntParam(1024, _is_greater_or_equal_0), + IntParam(1024, validate=lambda x: x >= 0), in_c_key=False, ) @@ -1023,16 +998,6 @@ def add_optimizer_configvars(): ) -def add_metaopt_configvars(): - config.add( - "metaopt__verbose", - "0 for silent, 1 for only warnings, 2 for full output with" - "timings and selected implementation", - IntParam(0), - in_c_key=False, - ) - - def add_vm_configvars(): config.add( "profile", @@ -1066,26 +1031,6 @@ def add_vm_configvars(): ) -def add_deprecated_configvars(): - # TODO: remove this? Agree - config.add( - "unittests__rseed", - "Seed to use for randomized unit tests. " - "Special value 'random' means using a seed of None.", - StrParam(666, validate=_good_seem_param), - in_c_key=False, - ) - - config.add( - "warn__round", - "Warn when using `tensor.round` with the default mode. " - "Round changed its default from `half_away_from_zero` to " - "`half_to_even` to have the same default as NumPy.", - BoolParam(_warn_default("0.9")), - in_c_key=False, - ) - - def add_scan_configvars(): config.add( "scan__allow_gc", @@ -1120,12 +1065,6 @@ def add_numba_configvars(): ) -def _default_compiledirname() -> str: - formatted = config.compiledir_format % _compiledir_format_dict - safe = re.sub(r"[\(\)\s,]+", "_", formatted) - return safe - - def _filter_base_compiledir(path: str | Path) -> Path: # Expand '~' in path return Path(path).expanduser() @@ -1198,29 +1137,42 @@ def _get_home_dir() -> Path: return Path(windowsfail_home) -_compiledir_format_dict = { - "platform": platform.platform(), - "processor": platform.processor(), - "python_version": platform.python_version(), - "python_bitwidth": LOCAL_BITWIDTH, - "python_int_bitwidth": PYTHON_INT_BITWIDTH, - "pytensor_version": pytensor.__version__, - "numpy_version": np.__version__, - "gxx_version": "xxx", - "hostname": platform.node(), -} - - -def _default_compiledir() -> Path: - return config.base_compiledir / _default_compiledirname() - - def add_caching_dir_configvars(): - _compiledir_format_dict["gxx_version"] = (gcc_version_str.replace(" ", "_"),) - _compiledir_format_dict["short_platform"] = short_platform() - # Allow to have easily one compiledir per device. - _compiledir_format_dict["device"] = config.device - compiledir_format_keys = ", ".join(sorted(_compiledir_format_dict)) + # TODO: Does anybody need to customize these keys ever? + compiledir_format_keys = { + "platformshort_platform", + "processor", + "python_version", + "python_bitwidth", + "python_int_bitwidth", + "pytensor_version", + "numpy_version", + "gxx_version", + "hostname", + } + + def _default_compiledir() -> Path: + from pytensor import __version__ + + compiledir_dict = { + "platform": platform.platform(), + "short_platform": short_platform(), + "processor": platform.processor(), + "python_version": platform.python_version(), + "python_bitwidth": LOCAL_BITWIDTH, + "python_int_bitwidth": PYTHON_INT_BITWIDTH, + "pytensor_version": __version__, + "numpy_version": np.__version__, + "gxx_version": (config.gcc_version_str.replace(" ", "_"),), + "hostname": platform.node(), + } + assert set(compiledir_dict) == compiledir_format_keys + formatted = config.compiledir_format % compiledir_dict + safe = re.sub(r"[\(\)\s,]+", "_", formatted) + + return config.base_compiledir / safe + + compiledir_format_keys = ", ".join(sorted(compiledir_format_keys)) _default_compiledir_format = ( "compiledir_%(short_platform)s-%(processor)s-" "%(python_version)s-%(python_bitwidth)s" @@ -1233,7 +1185,7 @@ def add_caching_dir_configvars(): f"""\ Format string for platform-dependent compiled module subdirectory (relative to base_compiledir). - Available keys: {compiledir_format_keys}. Defaults to {_default_compiledir_format}. + Available keys: {compiledir_format_keys}. """ ) ), @@ -1273,32 +1225,13 @@ def add_caching_dir_configvars(): # The functions below register config variables into the config instance above. add_basic_configvars() add_compile_configvars() +add_c_compile_configvars() add_tensor_configvars() add_traceback_configvars() add_error_and_warning_configvars() add_testvalue_and_checking_configvars() add_multiprocessing_configvars() add_optimizer_configvars() -# TODO: Module-specific configs should probably be added upon import of the module. -# This would mean either calling the function from there, or even moving all the related code there. -# Blas-related config are a special pain-point, because their addition depends on a lot of stuff from -# that module, which introduces a circular dependency! -add_metaopt_configvars() -add_deprecated_configvars() add_vm_configvars() add_numba_configvars() - -# TODO: `gcc_version_str` is used by other modules.. Should it become an immutable config var? -if config.cxx != "": - try: - p_out = output_subprocess_Popen([config.cxx, "-dumpversion"]) - gcc_version_str = p_out[0].strip().decode() - except OSError: - # Typically means gcc cannot be found. - gcc_version_str = "GCC_NOT_FOUND" -else: - gcc_version_str = "GCC_NOT_FOUND" - -# TODO: The caching dir resolution is a procedural mess of helper functions, local variables -# and config definitions. And the result is also not particularly pretty.. add_caching_dir_configvars() diff --git a/pytensor/configparser.py b/pytensor/configparser.py index 7de4c87c86..485797680a 100644 --- a/pytensor/configparser.py +++ b/pytensor/configparser.py @@ -70,7 +70,6 @@ class PyTensorConfigParser: warn_float64: str pickle_test_value: bool cast_policy: str - device: str print_global_stats: bool unpickle_function: bool # add_compile_configvars @@ -512,30 +511,6 @@ def _apply(self, value): ) -class DeviceParam(ConfigParam): - def __init__(self, default, *options, **kwargs): - super().__init__( - default, apply=self._apply, mutable=kwargs.get("mutable", True) - ) - - def _apply(self, val): - if val.startswith("opencl") or val.startswith("cuda") or val.startswith("gpu"): - raise ValueError( - "You are trying to use the old GPU back-end. " - "It was removed from PyTensor." - ) - elif val == self.default: - return val - raise ValueError( - f'Invalid value ("{val}") for configuration ' - f'variable "{self.name}". Valid options start with ' - 'one of "cpu".' - ) - - def __str__(self): - return f"{self.name} ({self.default})" - - def parse_config_string( config_string: str, issue_warnings: bool = True ) -> dict[str, str]: diff --git a/pytensor/graph/__init__.py b/pytensor/graph/__init__.py index 5753479d25..c1b024b36f 100644 --- a/pytensor/graph/__init__.py +++ b/pytensor/graph/__init__.py @@ -7,7 +7,7 @@ Constant, clone, ) -from pytensor.graph.traversal import ancestors, graph_inputs +from pytensor.graph.traversal import ancestors, graph_inputs, explicit_graph_inputs from pytensor.graph.replace import clone_replace, graph_replace, vectorize_graph from pytensor.graph.op import Op from pytensor.graph.type import Type diff --git a/pytensor/ifelse.py b/pytensor/ifelse.py index f8e033a431..7669990a12 100644 --- a/pytensor/ifelse.py +++ b/pytensor/ifelse.py @@ -17,8 +17,7 @@ import numpy as np -import pytensor.tensor as pt -from pytensor import as_symbolic +from pytensor.basic import as_symbolic from pytensor.compile import optdb from pytensor.configdefaults import config from pytensor.graph.basic import Apply, Variable @@ -27,7 +26,13 @@ from pytensor.graph.rewriting.basic import GraphRewriter, in2out, node_rewriter from pytensor.graph.traversal import apply_depends_on from pytensor.graph.type import HasDataType, HasShape +from pytensor.tensor import as_tensor_variable +from pytensor.tensor.basic import Alloc, zeros_like +from pytensor.tensor.blockwise import Blockwise +from pytensor.tensor.elemwise import DimShuffle, Elemwise +from pytensor.tensor.math import Argmax, Dot, Max from pytensor.tensor.shape import Reshape, Shape, SpecifyShape +from pytensor.tensor.subtensor import IncSubtensor, Subtensor if TYPE_CHECKING: @@ -160,7 +165,7 @@ def make_node(self, condition: "TensorLike", *true_false_branches: Any): f"{int(2 * self.n_outs)}, got {len(true_false_branches)}" ) - condition = pt.basic.as_tensor_variable(condition) + condition = as_tensor_variable(condition) if condition.type.ndim > 0: raise TypeError("The condition argument must be a truthy scalar value") @@ -262,14 +267,14 @@ def grad(self, ins, grads): [condition] + grads + [ - pt.basic.zeros_like(t, dtype=grads[i].dtype) + zeros_like(t, dtype=grads[i].dtype) for i, t in enumerate(inputs_true_branch) ] ) inputs_false_grad = ( [condition] + [ - pt.basic.zeros_like(f, dtype=grads[i].dtype) + zeros_like(f, dtype=grads[i].dtype) for i, f in enumerate(inputs_false_branch) ] + grads @@ -482,15 +487,15 @@ def cond_make_inplace(fgraph, node): Shape, SpecifyShape, Reshape, - pt.math.Dot, - pt.math.Max, - pt.math.Argmax, - pt.subtensor.Subtensor, - pt.subtensor.IncSubtensor, - pt.basic.Alloc, - pt.elemwise.Elemwise, - pt.elemwise.DimShuffle, - pt.blockwise.Blockwise, + Dot, + Max, + Argmax, + Subtensor, + IncSubtensor, + Alloc, + Elemwise, + DimShuffle, + Blockwise, ) diff --git a/pytensor/link/c/basic.py b/pytensor/link/c/basic.py index 7e1a779c2e..4873ca6879 100644 --- a/pytensor/link/c/basic.py +++ b/pytensor/link/c/basic.py @@ -1368,7 +1368,7 @@ def cmodule_key_( # DynamicModule always add the include sig.append(f"NPY_ABI_VERSION=0x{NDARRAY_C_VERSION:X}") if c_compiler: - sig.append("c_compiler_str=" + c_compiler.version_str()) + sig.append(f"c_compiler_str={config.cxx} {config.gcc_version_str}") # IMPORTANT: The 'md5' prefix is used to isolate the compilation # parameters from the rest of the key. If you want to add more key diff --git a/pytensor/link/c/cmodule.py b/pytensor/link/c/cmodule.py index f7389231ef..5004df93fe 100644 --- a/pytensor/link/c/cmodule.py +++ b/pytensor/link/c/cmodule.py @@ -30,7 +30,7 @@ # 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.configdefaults import config from pytensor.configparser import BoolParam, StrParam from pytensor.graph.op import Op from pytensor.utils import ( @@ -1817,10 +1817,6 @@ def std_lib_dirs(): return std_lib_dirs_and_libs()[1] -def gcc_version(): - return gcc_version_str - - @is_GCCLLVMType def gcc_llvm() -> bool | None: """ @@ -2062,10 +2058,6 @@ class GCC_compiler(Compiler): supports_amdlibm = True - @staticmethod - def version_str(): - return config.cxx + " " + gcc_version_str - @staticmethod def compile_args(march_flags=True): cxxflags = [flag for flag in config.gcc__cxxflags.split(" ") if flag] @@ -2266,7 +2258,7 @@ def join_options(init_part): # OK continue # Check the version of GCC - version = gcc_version_str.split(".") + version = config.gcc_version_str.split(".") if len(version) != 3: # Unexpected, but should not be a problem continue diff --git a/pytensor/scan/basic.py b/pytensor/scan/basic.py index 215c1d06a2..8781718c43 100644 --- a/pytensor/scan/basic.py +++ b/pytensor/scan/basic.py @@ -21,7 +21,6 @@ from pytensor.tensor.math import minimum from pytensor.tensor.shape import shape_padleft from pytensor.tensor.type import TensorType, integer_dtypes -from pytensor.updates import OrderedUpdates if typing.TYPE_CHECKING: @@ -1170,7 +1169,7 @@ def wrap_into_list(x): # and so on ... ## - update_map = OrderedUpdates() + update_map = {} def remove_dimensions(outs, offsets=None): out_ls = [] diff --git a/pytensor/sparse/basic.py b/pytensor/sparse/basic.py index 3250fa7ca0..9d8d0c80c2 100644 --- a/pytensor/sparse/basic.py +++ b/pytensor/sparse/basic.py @@ -15,8 +15,8 @@ from numpy.lib.stride_tricks import as_strided import pytensor -from pytensor import _as_symbolic, as_symbolic from pytensor import scalar as ps +from pytensor.basic import _as_symbolic, as_symbolic from pytensor.configdefaults import config from pytensor.gradient import DisconnectedType, disconnected_type, grad_undefined from pytensor.graph.basic import Apply, Constant, Variable diff --git a/pytensor/sparse/rewriting.py b/pytensor/sparse/rewriting.py index d992635298..5012cce9cb 100644 --- a/pytensor/sparse/rewriting.py +++ b/pytensor/sparse/rewriting.py @@ -186,7 +186,7 @@ def c_code_cache_version(self): @node_rewriter([spm.AddSD]) def local_inplace_addsd_ccode(fgraph, node): """Rewrite to insert inplace versions of `AddSD`.""" - if isinstance(node.op, spm.AddSD) and config.cxx: + if isinstance(node.op, spm.AddSD): out_dtype = ps.upcast(*[inp.type.dtype for inp in node.inputs]) if out_dtype != node.inputs[1].dtype: return @@ -225,7 +225,7 @@ def local_addsd_ccode(fgraph, node): Convert AddSD to faster AddSD_ccode. """ - if isinstance(node.op, spm.AddSD) and config.cxx: + if isinstance(node.op, spm.AddSD): new_node = AddSD_ccode(format=node.inputs[0].type.format)(*node.inputs) return [new_node] return False diff --git a/pytensor/tensor/__init__.py b/pytensor/tensor/__init__.py index 003964caeb..1a557fd20c 100644 --- a/pytensor/tensor/__init__.py +++ b/pytensor/tensor/__init__.py @@ -1,15 +1,14 @@ """Symbolic tensor types and constructor functions.""" -from collections.abc import Callable, Sequence +from collections.abc import Sequence from functools import singledispatch -from typing import TYPE_CHECKING, Any, NoReturn, Optional, Union +from typing import TYPE_CHECKING, Union -from pytensor.graph.basic import Constant, Variable -from pytensor.graph.op import Op +from pytensor.graph import Constant, Op, Variable if TYPE_CHECKING: - from numpy.typing import ArrayLike, NDArray + from numpy.typing import ArrayLike TensorLike = Union[Variable, Sequence[Variable], "ArrayLike"] @@ -119,10 +118,6 @@ def _get_vector_length_Constant(op: Op | Variable, var: Constant) -> int: from pytensor.tensor import signal from pytensor.tensor import optimize -# For backward compatibility -from pytensor.tensor import nlinalg -from pytensor.tensor import slinalg - # isort: on # Allow accessing numpy constants from pytensor.tensor from numpy import e, euler_gamma, inf, nan, newaxis, pi @@ -162,7 +157,5 @@ def _get_vector_length_Constant(op: Op | Variable, var: Constant) -> int: # isort: off from pytensor.tensor.einsum import einsum from pytensor.tensor.functional import vectorize +from pytensor.tensor import random # isort: on - - -__all__ = ["random"] # noqa: F405 diff --git a/pytensor/tensor/interpolate.py b/pytensor/tensor/interpolate.py index f598695784..0301eb181a 100644 --- a/pytensor/tensor/interpolate.py +++ b/pytensor/tensor/interpolate.py @@ -2,7 +2,7 @@ from difflib import get_close_matches from typing import Literal, get_args -from pytensor import Variable +from pytensor.graph import Variable from pytensor.tensor.basic import as_tensor_variable, switch from pytensor.tensor.extra_ops import searchsorted from pytensor.tensor.functional import vectorize diff --git a/pytensor/tensor/reshape.py b/pytensor/tensor/reshape.py index 98024f28ca..e485a5d8a4 100644 --- a/pytensor/tensor/reshape.py +++ b/pytensor/tensor/reshape.py @@ -5,10 +5,8 @@ import numpy as np from numpy.lib._array_utils_impl import normalize_axis_index, normalize_axis_tuple -from pytensor import Variable from pytensor.gradient import disconnected_type -from pytensor.graph import Apply -from pytensor.graph.op import Op +from pytensor.graph import Apply, Op, Variable from pytensor.graph.replace import _vectorize_node from pytensor.scalar import ScalarVariable from pytensor.tensor import TensorLike, as_tensor_variable diff --git a/pytensor/tensor/rewriting/elemwise.py b/pytensor/tensor/rewriting/elemwise.py index dc30beedf3..5e2c48abc6 100644 --- a/pytensor/tensor/rewriting/elemwise.py +++ b/pytensor/tensor/rewriting/elemwise.py @@ -9,7 +9,7 @@ from warnings import warn from pytensor.compile.function.types import Supervisor -from pytensor.compile.mode import get_target_language, optdb +from pytensor.compile.mode import optdb from pytensor.configdefaults import config from pytensor.graph.basic import Apply, Variable from pytensor.graph.destroyhandler import DestroyHandler, inplace_candidates @@ -518,6 +518,7 @@ def flatten_nested_add_mul(fgraph, node): def elemwise_max_operands_fct(node) -> int: # `Elemwise.perform` uses NumPy ufuncs and they are limited to 32 operands (inputs and outputs) + # FIXME: config.cxx is not a good criteria! if not config.cxx: return 32 return 1024 @@ -1001,10 +1002,6 @@ def local_careduce_fusion(fgraph, node): if len(fgraph.clients[elm_outputs[0]]) > 1: return False - # Don't form the fusion when the target language is Python - if get_target_language() == ("py",): - return False - if not elm_scalar_op.supports_c_code(elm_inputs, elm_outputs): return None @@ -1201,6 +1198,7 @@ def constant_fold_branches_of_add_mul(fgraph, node): dfs_rewriter(local_careduce_fusion), "fast_run", "fusion", + "cxx_only", position=10, ) fuse_seqopt.register( diff --git a/pytensor/tensor/rewriting/linalg.py b/pytensor/tensor/rewriting/linalg.py index 2c17020cd9..0c7ab7025a 100644 --- a/pytensor/tensor/rewriting/linalg.py +++ b/pytensor/tensor/rewriting/linalg.py @@ -4,9 +4,8 @@ import numpy as np -from pytensor import Variable from pytensor import tensor as pt -from pytensor.graph import Apply, FunctionGraph +from pytensor.graph import Apply, FunctionGraph, Variable from pytensor.graph.rewriting.basic import ( copy_stack_trace, node_rewriter, diff --git a/pytensor/tensor/rewriting/ofg.py b/pytensor/tensor/rewriting/ofg.py index 098d380fad..f7ae945952 100644 --- a/pytensor/tensor/rewriting/ofg.py +++ b/pytensor/tensor/rewriting/ofg.py @@ -1,9 +1,8 @@ from typing import cast -from pytensor import Variable, clone_replace from pytensor.compile import optdb from pytensor.compile.builders import OpFromGraph -from pytensor.graph import Apply, node_rewriter +from pytensor.graph import Apply, Variable, clone_replace, node_rewriter from pytensor.graph.rewriting.basic import copy_stack_trace, dfs_rewriter from pytensor.tensor.basic import AllocDiag from pytensor.tensor.rewriting.basic import register_specialize diff --git a/pytensor/tensor/rewriting/subtensor_lift.py b/pytensor/tensor/rewriting/subtensor_lift.py index b21ad516ab..a415e6b725 100644 --- a/pytensor/tensor/rewriting/subtensor_lift.py +++ b/pytensor/tensor/rewriting/subtensor_lift.py @@ -4,9 +4,14 @@ import numpy as np from numpy.lib.array_utils import normalize_axis_index, normalize_axis_tuple -from pytensor import Variable from pytensor.compile import optdb -from pytensor.graph import Constant, FunctionGraph, node_rewriter, vectorize_graph +from pytensor.graph import ( + Constant, + FunctionGraph, + Variable, + node_rewriter, + vectorize_graph, +) from pytensor.graph.rewriting.basic import NodeRewriter, copy_stack_trace from pytensor.scalar import basic as ps from pytensor.tensor.basic import ( diff --git a/pytensor/tensor/type_other.py b/pytensor/tensor/type_other.py index a60563f9b3..5aeba66aa4 100644 --- a/pytensor/tensor/type_other.py +++ b/pytensor/tensor/type_other.py @@ -5,7 +5,7 @@ import numpy as np import pytensor -from pytensor import _as_symbolic +from pytensor.basic import _as_symbolic from pytensor.gradient import disconnected_type from pytensor.graph.basic import Apply, Constant, Variable from pytensor.graph.op import Op diff --git a/pytensor/updates.py b/pytensor/updates.py deleted file mode 100644 index fa1320200c..0000000000 --- a/pytensor/updates.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Defines Updates object for storing a (SharedVariable, new_value) mapping.""" - -import logging - -from pytensor.compile.sharedvalue import SharedVariable - - -__docformat__ = "restructuredtext en" - -logger = logging.getLogger("pytensor.updates") - - -# Relies on the fact that dict is ordered, otherwise updates will be applied -# in a non-deterministic order. -class OrderedUpdates(dict): - """ - Dict-like mapping from SharedVariable keys to their new values. - - This mapping supports the use of the "+" operator for the union of updates. - """ - - def __init__(self, *key, **kwargs): - super().__init__(*key, **kwargs) - for key in self: - if not isinstance(key, SharedVariable): - raise TypeError( - "OrderedUpdates keys must inherit from SharedVariable", key - ) - - def __setitem__(self, key, value): - if isinstance(key, SharedVariable): - # TODO: consider doing error-checking on value. - # insist that it is an PyTensor variable? Have the right type? - # This could have weird consequences - for example a - - return super().__setitem__(key, value) - else: - raise TypeError("OrderedUpdates keys must inherit from SharedVariable", key) - - def update(self, other=None): - if other is None: - return - for key, val in dict(other).items(): - if key in self: - if self[key] == val: - continue - raise KeyError("Collision", key) - self[key] = val # __setitem__ does type-checking - - def __add__(self, other): - rval = OrderedUpdates() - rval.update(self) - rval.update(other) - return rval - - def __radd__(other, self): - rval = OrderedUpdates() - rval.update(other) - rval.update(self) - return rval diff --git a/pytensor/xtensor/rewriting/indexing.py b/pytensor/xtensor/rewriting/indexing.py index 25a0f80dd4..35ebcf2ef3 100644 --- a/pytensor/xtensor/rewriting/indexing.py +++ b/pytensor/xtensor/rewriting/indexing.py @@ -1,6 +1,6 @@ from itertools import zip_longest -from pytensor import as_symbolic +from pytensor.basic import as_symbolic from pytensor.graph import Constant, node_rewriter from pytensor.tensor import TensorType, arange, specify_shape from pytensor.tensor.subtensor import _non_consecutive_adv_indexing, inc_subtensor diff --git a/pytensor/xtensor/type.py b/pytensor/xtensor/type.py index db6a66036d..9dc5ed6a13 100644 --- a/pytensor/xtensor/type.py +++ b/pytensor/xtensor/type.py @@ -30,7 +30,8 @@ import numpy as np import pytensor.xtensor as px -from pytensor import _as_symbolic, config +from pytensor import config +from pytensor.basic import _as_symbolic from pytensor.graph import Apply, Constant from pytensor.graph.basic import OptionalApplyType, Variable from pytensor.graph.type import HasDataType, HasShape, Type diff --git a/tests/compile/test_mode.py b/tests/compile/test_mode.py index 12b2068aaa..2d46a92cc1 100644 --- a/tests/compile/test_mode.py +++ b/tests/compile/test_mode.py @@ -1,18 +1,14 @@ import copy -import pytest - from pytensor.compile.function import function from pytensor.compile.mode import ( AddFeatureOptimizer, Mode, get_default_mode, - get_target_language, ) from pytensor.configdefaults import config from pytensor.graph.features import NoOutputFromInplace from pytensor.graph.rewriting.db import RewriteDatabaseQuery, SequenceDB -from pytensor.link.basic import LocalLinker from pytensor.link.jax import JAXLinker from pytensor.tensor.math import dot, tanh from pytensor.tensor.type import matrix, vector @@ -125,41 +121,6 @@ def test_modes(self): assert not hasattr(linker, "fgraph") or linker.fgraph is None -def test_get_target_language(): - with config.change_flags(mode=Mode(linker="py")): - res = get_target_language() - assert res == ("py",) - - res = get_target_language(Mode(linker="py")) - assert res == ("py",) - - res = get_target_language(Mode(linker="c")) - assert res == ("c",) - - res = get_target_language(Mode(linker="c|py")) - assert res == ("c", "py") - - res = get_target_language(Mode(linker="vm")) - assert res == ("c", "py") - - with config.change_flags(cxx=""): - res = get_target_language(Mode(linker="vm")) - assert res == ("py",) - - res = get_target_language(Mode(linker="jax")) - assert res == ("jax",) - - res = get_target_language(Mode(linker="numba")) - assert res == ("numba",) - - class MyLinker(LocalLinker): - pass - - test_mode = Mode(linker=MyLinker()) - with pytest.raises(Exception): - get_target_language(test_mode) - - def test_predefined_modes_respected(): default_mode = get_default_mode() assert not isinstance(default_mode.linker, JAXLinker) diff --git a/tests/link/numba/test_subtensor.py b/tests/link/numba/test_subtensor.py index b700172779..af723e91e9 100644 --- a/tests/link/numba/test_subtensor.py +++ b/tests/link/numba/test_subtensor.py @@ -5,7 +5,8 @@ import pytensor.scalar as ps import pytensor.tensor as pt -from pytensor import Mode, as_symbolic +from pytensor import Mode +from pytensor.basic import as_symbolic from pytensor.tensor import as_tensor from pytensor.tensor.subtensor import ( AdvancedIncSubtensor, diff --git a/tests/tensor/test_subtensor.py b/tests/tensor/test_subtensor.py index 6f79694e25..8655837838 100644 --- a/tests/tensor/test_subtensor.py +++ b/tests/tensor/test_subtensor.py @@ -11,7 +11,7 @@ import pytensor import pytensor.scalar as scal import pytensor.tensor.basic as ptb -from pytensor import function +from pytensor import basic, function from pytensor.compile import DeepCopyOp, shared from pytensor.compile.io import In from pytensor.compile.mode import Mode, get_default_mode @@ -2798,7 +2798,7 @@ def test_AdvancedSubtensor_bool_mixed(self): def test_advanced_subtensor_constant_slice(self): x = dmatrix("x") - constant_slice = pytensor.as_symbolic(slice(1, None, None)) + constant_slice = basic.as_symbolic(slice(1, None, None)) assert isinstance(constant_slice, Constant) adv_indices = ptb.constant(np.zeros((2, 3)), dtype="int") y = advanced_subtensor(x, constant_slice, adv_indices) diff --git a/tests/tensor/test_type_other.py b/tests/tensor/test_type_other.py index 0d9131516d..7728416095 100644 --- a/tests/tensor/test_type_other.py +++ b/tests/tensor/test_type_other.py @@ -1,7 +1,7 @@ """This file don't test everything. It only test one past crash error.""" import pytensor -from pytensor import as_symbolic +from pytensor.basic import as_symbolic from pytensor.graph.basic import Constant from pytensor.tensor.math import argmax from pytensor.tensor.type import iscalar, vector diff --git a/tests/test_basic.py b/tests/test_basic.py new file mode 100644 index 0000000000..bc058d8b1a --- /dev/null +++ b/tests/test_basic.py @@ -0,0 +1,41 @@ +import pytensor + + +def test_root_module_not_polluted(): + module_items = sorted(i for i in dir(pytensor) if not i.startswith("__")) + assert module_items == [ + "In", + "Lop", + "Mode", + "OpFromGraph", + "Out", + "Rop", + "basic", + "compile", + "config", + "configdefaults", + "configparser", + "dprint", + "foldl", + "foldr", + "function", + "grad", + "gradient", + "graph", + "ifelse", + "link", + "map", + "misc", + "npy_2_compat", + "printing", + "raise_op", + "reduce", + "scalar", + "scan", + "shared", + "sparse", + "tensor", + "utils", + "wrap_jax", + "wrap_py", + ] diff --git a/tests/test_config.py b/tests/test_config.py index ea1de66561..a89fcdae4e 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -143,19 +143,6 @@ def test_enumstr(self): with pytest.raises(ValueError, match="Non-str value"): configparser.EnumStr(default="red", options=["red", 12, "yellow"]) - def test_deviceparam(self): - cp = configparser.DeviceParam("cpu", mutable=False) - assert cp.default == "cpu" - with pytest.raises(ValueError, match="It was removed from PyTensor"): - cp._apply("cuda123") - with pytest.raises(ValueError, match="It was removed from PyTensor"): - cp._apply("gpu123") - with pytest.raises( - ValueError, match='Valid options start with one of "cpu"\\.' - ): - cp._apply("notadevice") - assert str(cp) == "unnamed (cpu)" - def test_config_context(): root = _create_test_config() diff --git a/tests/test_updates.py b/tests/test_updates.py deleted file mode 100644 index 73f6394675..0000000000 --- a/tests/test_updates.py +++ /dev/null @@ -1,70 +0,0 @@ -import pytest - -import pytensor -from pytensor.tensor.type import vector -from pytensor.updates import OrderedUpdates - - -class TestUpdates: - def test_updates_init(self): - with pytest.raises(TypeError): - OrderedUpdates(dict(d=3)) - - sv = pytensor.shared("asdf") - # TODO FIXME: Not a real test. - OrderedUpdates({sv: 3}) - - def test_updates_setitem(self): - up = OrderedUpdates() - - # keys have to be SharedVariables - with pytest.raises(TypeError): - up.__setitem__(5, 7) - with pytest.raises(TypeError): - up.__setitem__(vector(), 7) - - # TODO FIXME: Not a real test. - up[pytensor.shared(88)] = 7 - - def test_updates_add(self): - up1 = OrderedUpdates() - up2 = OrderedUpdates() - - a = pytensor.shared("a") - b = pytensor.shared("b") - - assert not up1 + up2 - - up1[a] = 5 - - # test that addition works - assert up1 - assert up1 + up2 - assert not up2 - - assert len(up1 + up2) == 1 - assert (up1 + up2)[a] == 5 - - up2[b] = 7 - assert up1 - assert up1 + up2 - assert up2 - - assert len(up1 + up2) == 2 - assert (up1 + up2)[a] == 5 - assert (up1 + up2)[b] == 7 - - assert a in (up1 + up2) - assert b in (up1 + up2) - - # this works even though there is a collision - # because values all match - assert len(up1 + up1 + up1) == 1 - - up2[a] = 8 # a gets different value in up1 and up2 - with pytest.raises(KeyError): - up1 + up2 - - # TODO FIXME: Not a real test. - # reassigning to a key works fine right? - up2[a] = 10