pytensor 2.36.0__cp314-cp314-win_amd64.whl

pytensor/__init__.py

@@ -0,0 +1,176 @@
+"""
+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"]
+
+
+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.compile import (
+    In,
+    Mode,
+    Out,
+    ProfileStats,
+    predefined_linkers,
+    predefined_modes,
+    predefined_optimizers,
+    shared,
+)
+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.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
+from pytensor.link.jax.ops import wrap_jax
+# isort: on
+
+
+# Some config variables are registered by submodules. Only after all those
+# imports were executed, we can warn about remaining flags provided by the user
+# through PYTENSOR_FLAGS.
+config.warn_unused_flags()

pytensor/_version.py

@@ -0,0 +1,21 @@
+
+# This file was generated by 'versioneer.py' (0.29) from
+# revision-control system data, or from the parent directory name of an
+# unpacked source archive. Distribution tarballs contain a pre-generated copy
+# of this file.
+
+import json
+
+version_json = '''
+{
+ "date": "2025-12-15T22:58:11+0100",
+ "dirty": false,
+ "error": null,
+ "full-revisionid": "6319fac84ebc71ad079c5986c6c302a52e254198",
+ "version": "2.36.0"
+}
+'''  # END VERSION_JSON
+
+
+def get_versions():
+    return json.loads(version_json)

pytensor/bin/pytensor_cache.py

@@ -0,0 +1,112 @@
+import logging
+import os
+import shutil
+import sys
+from pathlib import Path
+
+
+if sys.platform == "win32":
+    config_for_pytensor_cache_script = "cxx=,device=cpu"
+    pytensor_flags = (
+        os.environ["PYTENSOR_FLAGS"] if "PYTENSOR_FLAGS" in os.environ else ""
+    )
+    if pytensor_flags:
+        pytensor_flags += ","
+    pytensor_flags += config_for_pytensor_cache_script
+    os.environ["PYTENSOR_FLAGS"] = pytensor_flags
+
+import pytensor
+import pytensor.compile.compiledir
+from pytensor import config
+from pytensor.link.c.basic import get_module_cache
+
+
+_logger = logging.getLogger("pytensor.bin.pytensor-cache")
+
+
+def print_help(exit_status):
+    if exit_status:
+        print(f'command "{" ".join(sys.argv)}" not recognized')
+    print('Type "pytensor-cache" to print the cache location')
+    print('Type "pytensor-cache help" to print this help')
+    print('Type "pytensor-cache clear" to erase the cache')
+    print('Type "pytensor-cache list" to print the cache content')
+    print('Type "pytensor-cache unlock" to unlock the cache directory')
+    print('Type "pytensor-cache cleanup" to delete keys in the old format/code version')
+    print('Type "pytensor-cache purge" to force deletion of the cache directory')
+    print(
+        'Type "pytensor-cache basecompiledir" '
+        "to print the parent of the cache directory"
+    )
+    print(
+        'Type "pytensor-cache basecompiledir list" '
+        "to print the content of the base compile dir"
+    )
+    print(
+        'Type "pytensor-cache basecompiledir purge" '
+        "to remove everything in the base compile dir, "
+        "that is, erase ALL cache directories"
+    )
+    sys.exit(exit_status)
+
+
+def main():
+    if len(sys.argv) == 1:
+        print(config.compiledir)
+    elif len(sys.argv) == 2:
+        if sys.argv[1] == "help":
+            print_help(exit_status=0)
+        if sys.argv[1] == "clear":
+            # We skip the refresh on module cache creation because the refresh will
+            # be done when calling clear afterwards.
+            cache = get_module_cache(init_args=dict(do_refresh=False))
+            cache.clear(
+                unversioned_min_age=-1, clear_base_files=True, delete_if_problem=True
+            )
+
+            # Print a warning if some cached modules were not removed, so that the
+            # user knows he should manually delete them, or call
+            # pytensor-cache purge, # to properly clear the cache.
+            items = list(Path(cache.dirname).glob("tmp*"))
+            if items:
+                _logger.warning(
+                    "There remain elements in the cache dir that you may "
+                    f"need to erase manually. The cache dir is:\n  {config.compiledir}\n"
+                    'You can also call "pytensor-cache purge" to '
+                    "remove everything from that directory."
+                )
+                _logger.debug(f"Remaining elements ({len(items)}): {items}")
+            numba_cache_dir: Path = config.base_compiledir / "numba"
+            shutil.rmtree(numba_cache_dir, ignore_errors=True)
+
+        elif sys.argv[1] == "list":
+            pytensor.compile.compiledir.print_compiledir_content()
+        elif sys.argv[1] == "cleanup":
+            pytensor.compile.compiledir.cleanup()
+            cache = get_module_cache(init_args=dict(do_refresh=False))
+            cache.clear_old()
+        elif sys.argv[1] == "unlock":
+            pytensor.compile.compilelock.force_unlock(config.compiledir)
+            print("Lock successfully removed!")
+        elif sys.argv[1] == "purge":
+            pytensor.compile.compiledir.compiledir_purge()
+            numba_cache_dir: Path = config.base_compiledir / "numba"
+            shutil.rmtree(numba_cache_dir, ignore_errors=True)
+        elif sys.argv[1] == "basecompiledir":
+            # Simply print the base_compiledir
+            print(pytensor.config.base_compiledir)
+        else:
+            print_help(exit_status=1)
+    elif len(sys.argv) == 3 and sys.argv[1] == "basecompiledir":
+        if sys.argv[2] == "list":
+            pytensor.compile.compiledir.basecompiledir_ls()
+        elif sys.argv[2] == "purge":
+            pytensor.compile.compiledir.basecompiledir_purge()
+        else:
+            print_help(exit_status=1)
+    else:
+        print_help(exit_status=1)
+
+
+if __name__ == "__main__":
+    main()

pytensor/breakpoint.py

@@ -0,0 +1,161 @@
+import numpy as np
+
+from pytensor.gradient import DisconnectedType
+from pytensor.graph.basic import Apply, Variable
+from pytensor.graph.op import Op
+from pytensor.tensor.basic import as_tensor_variable
+
+
+class PdbBreakpoint(Op):
+    """
+    This is an identity-like op with the side effect of enforcing a
+    conditional breakpoint, inside an PyTensor function, based on a symbolic
+    scalar condition. It automatically detects available debuggers and uses
+    the first available in the following order: `pudb`, `ipdb`, or `pdb`.
+
+    :type name: String
+    :param name: name of the conditional breakpoint. To be printed when the
+                 breakpoint is activated.
+
+    :note: WARNING. At least one of the outputs of the op must be used
+                    otherwise the op will be removed from the PyTensor graph
+                    due to its outputs being unused
+
+    :note: WARNING. Employing the function inside an PyTensor graph can prevent
+                    PyTensor from applying certain optimizations to improve
+                    performance, reduce memory consumption and/or reduce
+                    numerical instability.
+
+            Detailed explanation:
+            As of 2014-12-01 the PdbBreakpoint op is not known by any
+            optimization. Setting a PdbBreakpoint op in the middle of a
+            pattern that is usually optimized out will block the optimization.
+
+    Example:
+
+    .. code-block:: python
+
+        import pytensor
+        import pytensor.tensor as pt
+        from pytensor.breakpoint import PdbBreakpoint
+
+        input = pt.fvector()
+        target = pt.fvector()
+
+        # Mean squared error between input and target
+        mse = (input - target) ** 2
+
+        # Conditional breakpoint to be activated if the total MSE is higher
+        # than 100. The breakpoint will monitor the inputs, targets as well
+        # as the individual error values
+        breakpointOp = PdbBreakpoint("MSE too high")
+        condition = pt.gt(mse.sum(), 100)
+        mse, monitored_input, monitored_target = breakpointOp(condition, mse,
+                                                              input, target)
+
+        # Compile the pytensor function
+        fct = pytensor.function([input, target], mse)
+
+        # Use the function
+        print fct([10, 0], [10, 5]) # Will NOT activate the breakpoint
+        print fct([0, 0], [10, 5]) # Will activate the breakpoint
+
+
+    """
+
+    __props__ = ("name",)
+
+    def __init__(self, name):
+        self.name = name
+
+    def make_node(self, condition, *monitored_vars):
+        # Ensure that condition is an PyTensor tensor
+        if not isinstance(condition, Variable):
+            condition = as_tensor_variable(condition)
+
+        # Validate that the condition is a scalar (else it is not obvious how
+        # is should be evaluated)
+        assert condition.ndim == 0
+
+        # Because the user might be tempted to instantiate PdbBreakpoint only
+        # once and apply it many times on different number of inputs, we must
+        # create a new instance of the op here, define the instance attributes
+        # (view_map and var_types) in that instance and then apply it on the
+        # inputs.
+        new_op = PdbBreakpoint(name=self.name)
+        new_op.view_map = {}
+        new_op.inp_types = []
+        for i in range(len(monitored_vars)):
+            # Every output i is a view of the input i+1 because of the input
+            # condition.
+            new_op.view_map[i] = [i + 1]
+            new_op.inp_types.append(monitored_vars[i].type)
+
+        # Build the Apply node
+        inputs = [condition, *monitored_vars]
+        outputs = [inp.type() for inp in monitored_vars]
+        return Apply(op=new_op, inputs=inputs, outputs=outputs)
+
+    def perform(self, node, inputs, output_storage):
+        condition = inputs[0]
+
+        if condition:
+            try:
+                monitored = [np.asarray(inp) for inp in inputs[1:]]
+            except Exception:
+                raise ValueError(
+                    "Some of the inputs to the PdbBreakpoint op "
+                    f"'{self.name}' could not be casted to NumPy arrays"
+                )
+
+            print("\n")  # noqa: T201
+            print("-------------------------------------------------")  # noqa: T201
+            print(f"Conditional breakpoint '{self.name}' activated\n")  # noqa: T201
+            print("The monitored variables are stored, in order,")  # noqa: T201
+            print("in the list variable 'monitored' as NumPy arrays.\n")  # noqa: T201
+            print("Their contents can be altered and, when execution")  # noqa: T201
+            print("resumes, the updated values will be used.")  # noqa: T201
+            print("-------------------------------------------------")  # noqa: T201
+
+            try:
+                import pudb
+
+                pudb.set_trace()
+            except ImportError:
+                try:
+                    import ipdb
+
+                    ipdb.set_trace()
+                except ImportError:
+                    import pdb
+
+                    pdb.set_trace()
+
+            # Take the new values in monitored, cast them back to their
+            # original type and store them in the output_storage
+            for i in range(len(output_storage)):
+                output_storage[i][0] = self.inp_types[i].filter(monitored[i])
+
+        else:
+            # Simply return views on the monitored variables
+            for i in range(len(output_storage)):
+                output_storage[i][0] = inputs[i + 1]
+
+    def grad(self, inputs, output_gradients):
+        return [DisconnectedType()(), *output_gradients]
+
+    def infer_shape(self, fgraph, inputs, input_shapes):
+        # Return the shape of every input but the condition (first input)
+        return input_shapes[1:]
+
+    def connection_pattern(self, node):
+        nb_inp = len(node.inputs)
+        nb_out = nb_inp - 1
+
+        # First input is connected to no output and every other input n is
+        # connected to input n-1
+        connections = [
+            [out_idx == inp_idx - 1 for out_idx in range(nb_out)]
+            for inp_idx in range(nb_inp)
+        ]
+        return connections

pytensor/compile/__init__.py

@@ -0,0 +1,63 @@
+from pytensor.compile.function.pfunc import pfunc, rebuild_collect_shared
+from pytensor.compile.function.types import (
+    AliasedMemoryError,
+    Function,
+    FunctionMaker,
+    Supervisor,
+    UnusedInputError,
+    alias_root,
+    convert_function_input,
+    fgraph_updated_vars,
+    get_info_on_inputs,
+    infer_reuse_pattern,
+    insert_deepcopy,
+    orig_function,
+    std_fgraph,
+    view_tree_set,
+)
+from pytensor.compile.io import In, Out, SymbolicInput, SymbolicOutput
+from pytensor.compile.mode import (
+    CVM,
+    FAST_COMPILE,
+    JAX,
+    NUMBA,
+    OPT_FAST_COMPILE,
+    OPT_FAST_RUN,
+    OPT_FAST_RUN_STABLE,
+    OPT_MERGE,
+    OPT_NONE,
+    OPT_O2,
+    OPT_O3,
+    OPT_STABILIZE,
+    OPT_UNSAFE,
+    PYTORCH,
+    AddDestroyHandler,
+    AddFeatureOptimizer,
+    C,
+    Mode,
+    PrintCurrentFunctionGraph,
+    get_default_mode,
+    get_mode,
+    local_useless,
+    optdb,
+    predefined_linkers,
+    predefined_modes,
+    predefined_optimizers,
+    register_linker,
+    register_mode,
+    register_optimizer,
+)
+from pytensor.compile.monitormode import MonitorMode
+from pytensor.compile.ops import (
+    DeepCopyOp,
+    FromFunctionOp,
+    ViewOp,
+    as_op,
+    deep_copy_op,
+    register_deep_copy_op_c_code,
+    register_view_op_c_code,
+    view_op,
+    wrap_py,
+)
+from pytensor.compile.profiling import ProfileStats
+from pytensor.compile.sharedvalue import SharedVariable, shared, shared_constructor