pytensor 2.24.2__tar.gz → 2.25.0__tar.gz

{pytensor-2.24.2/pytensor.egg-info → pytensor-2.25.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytensor
-Version: 2.24.2
+Version: 2.25.0
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License: .. _license:
@@ -75,7 +75,7 @@ License-File: LICENSE.txt
 Requires-Dist: setuptools>=59.0.0
 Requires-Dist: scipy<1.14,>=0.14
 Requires-Dist: numpy<2,>=1.17.0
-Requires-Dist: filelock
+Requires-Dist: filelock>=3.15
 Requires-Dist: etuples
 Requires-Dist: logical-unification
 Requires-Dist: miniKanren

{pytensor-2.24.2 → pytensor-2.25.0}/doc/conf.py

@@ -12,11 +12,12 @@
 # serve to show the default value.
 
 # If your extensions are in another directory, add it here. If the directory
-# is relative to the documentation root, use os.path.abspath to make it
+# is relative to the documentation root, use Path.absolute to make it
 # absolute, like shown here.
-# sys.path.append(os.path.abspath('some/directory'))
+# sys.path.append(str(Path("some/directory").absolute()))
 
 import os
+import inspect
 import sys
 import pytensor
 
@@ -236,11 +237,9 @@ def linkcode_resolve(domain, info):
         obj = sys.modules[info["module"]]
         for part in info["fullname"].split("."):
             obj = getattr(obj, part)
-        import inspect
-        import os
 
-        fn = inspect.getsourcefile(obj)
-        fn = os.path.relpath(fn, start=os.path.dirname(pytensor.__file__))
+        fn = Path(inspect.getsourcefile(obj))
+        fn = fn.relative_to(Path(__file__).parent)
         source, lineno = inspect.getsourcelines(obj)
         return fn, lineno, lineno + len(source) - 1
 

{pytensor-2.24.2 → pytensor-2.25.0}/doc/library/config.rst

@@ -510,116 +510,12 @@ import ``pytensor`` and print the config variable, as in:
 
     Removing these asserts can speed up execution.
 
-.. attribute:: config.dnn__enabled
-
-    String value: ``'auto'``, ``'True'``, ``'False'``, ``'no_check'``
-
-    Default: ``'auto'``
-
-    If ``'auto'``, automatically detect and use
-    `cuDNN <https://developer.nvidia.com/cudnn>`_ when it is available.
-    If cuDNN is unavailable, do not raise an error.
-
-    If ``'True'``, require the use of cuDNN.  If cuDNN is unavailable, raise an error.
-
-    If ``'False'``, neither use cuDNN nor check if it is available.
-
-    If ``'no_check'``, assume cuDNN is present and that the versions between the
-    header and library match.
-
-.. attribute:: config.dnn__include_path
-
-    Default: ``include`` sub-folder in CUDA root directory, or headers paths defined for the compiler.
-
-    Location of the cuDNN header.
-
-.. attribute:: config.dnn__library_path
-
-    Default: Library sub-folder (``lib64`` on Linux) in CUDA root directory, or
-    libraries paths defined for the compiler.
-
-    Location of the cuDNN library.
-
 .. attribute:: config.conv__assert_shape
 
     If ``True``, ``AbstractConv*`` :class:`Op`\s will verify that user-provided shapes
     match the run-time shapes. This is a debugging option, and may slow down
     compilation.
 
-.. attribute:: config.dnn.conv.workmem
-
-    Deprecated, use :attr:`config.dnn__conv__algo_fwd`.
-
-
-.. attribute:: config.dnn.conv.workmem_bwd
-
-    Deprecated, use :attr:`config.dnn__conv__algo_bwd_filter` and
-    :attr:`config.dnn__conv__algo_bwd_data` instead.
-
-.. attribute:: config.dnn__conv__algo_fwd
-
-    String value:
-    ``'small'``, ``'none'``, ``'large'``, ``'fft'``, ``'fft_tiling'``,
-    ``'winograd'``, ``'winograd_non_fused'``, ``'guess_once'``, ``'guess_on_shape_change'``,
-    ``'time_once'``, ``'time_on_shape_change'``.
-
-    Default: ``'small'``
-
-    3d convolution only support ``'none'``, ``'small'``, ``'fft_tiling'``, ``'guess_once'``,
-    ``'guess_on_shape_change'``, ``'time_once'``, ``'time_on_shape_change'``.
-
-.. attribute:: config.dnn.conv.algo_bwd
-
-    Deprecated, use :attr:`config.dnn__conv__algo_bwd_filter` and
-    :attr:`config.dnn__conv__algo_bwd_data` instead.
-
-.. attribute:: config.dnn__conv__algo_bwd_filter
-
-    String value:
-    ``'none'``, ``'deterministic'``, ``'fft'``, ``'small'``, ``'winograd_non_fused'``, ``'fft_tiling'``, ``'guess_once'``,
-    ``'guess_on_shape_change'``, ``'time_once'``, ``'time_on_shape_change'``.
-
-    Default: ``'none'``
-
-    3d convolution only supports ``'none'``, ``'small'``, ``'guess_once'``,
-    ``'guess_on_shape_change'``, ``'time_once'``, ``'time_on_shape_change'``.
-
-.. attribute:: config.dnn__conv__algo_bwd_data
-
-    String value:
-    ``'none'``, ``'deterministic'``, ``'fft'``, ``'fft_tiling'``, ``'winograd'``,
-    ``'winograd_non_fused'``, ``'guess_once'``, ``'guess_on_shape_change'``, ``'time_once'``,
-    ``'time_on_shape_change'``.
-
-    Default: ``'none'``
-
-    3d convolution only supports ``'none'``, ``'deterministic'``, ``'fft_tiling'``
-    ``'guess_once'``, ``'guess_on_shape_change'``, ``'time_once'``,
-    ``'time_on_shape_change'``.
-
-.. attribute:: config.magma__enabled
-
-    String value: ``'True'``, ``'False'``
-
-    Default: ``'False'``
-
-    If ``'True'``, use `magma <http://icl.cs.utk.edu/magma/>`_ for matrix
-    computations.
-
-    If ``'False'``, disable magma.
-
-.. attribute:: config.magma__include_path
-
-    Default: ``''``
-
-    Location of the magma headers.
-
-.. attribute:: config.magma__library_path
-
-    Default: ``''``
-
-    Location of the magma library.
-
 .. attribute:: config.ctc__root
 
     Default: ``''``

{pytensor-2.24.2 → pytensor-2.25.0}/doc/library/d3viz/index.rst

@@ -76,10 +76,9 @@ visualize it with :py:func:`pytensor.printing.pydotprint` as follows:
 .. code:: python
 
     from pytensor.printing import pydotprint
-    import os
+    from pathlib import Path
 
-    if not os.path.exists('examples'):
-        os.makedirs('examples')
+    Path("examples").mkdir(exist_ok=True)
     pydotprint(predict, 'examples/mlp.png')
 
 

{pytensor-2.24.2 → pytensor-2.25.0}/doc/troubleshooting.rst

@@ -259,7 +259,7 @@ PyTensor/BLAS speed test:
 
 .. code-block:: bash
 
-    python `python -c "import os, pytensor; print(os.path.dirname(pytensor.__file__))"`/misc/check_blas.py
+    python $(python -c "import pathlib, pytensor; print(pathlib.Path(pytensor.__file__).parent / 'misc/check_blas.py')")
 
 This will print a table with different versions of BLAS/numbers of
 threads on multiple CPUs. It will also print some PyTensor/NumPy

{pytensor-2.24.2 → pytensor-2.25.0}/pyproject.toml

@@ -49,7 +49,7 @@ dependencies = [
     "setuptools>=59.0.0",
     "scipy>=0.14,<1.14",
     "numpy>=1.17.0,<2",
-    "filelock",
+    "filelock>=3.15",
     "etuples",
     "logical-unification",
     "miniKanren",

{pytensor-2.24.2 → pytensor-2.25.0}/pytensor/__init__.py

@@ -23,9 +23,9 @@ __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 os
 import sys
 from functools import singledispatch
+from pathlib import Path
 from typing import Any, NoReturn, Optional
 
 from pytensor.version import version as __version__
@@ -52,10 +52,8 @@ def disable_log_handler(logger=pytensor_logger, handler=logging_default_handler)
 
 # Raise a meaningful warning/error if the pytensor directory is in the Python
 # path.
-rpath = os.path.realpath(__path__[0])
-for p in sys.path:
-    if os.path.realpath(p) != rpath:
-        continue
+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

{pytensor-2.24.2 → pytensor-2.25.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-07-08T14:11:01+0200",
+ "date": "2024-07-09T17:06:59+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "7f623fefbed89442737fd6d8ba11c32d2eab65ed",
- "version": "2.24.2"
+ "full-revisionid": "a6e79f28775c39996fb746bf3f369b422c94cfcf",
+ "version": "2.25.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.24.2 → pytensor-2.25.0}/pytensor/compile/compiledir.py

@@ -4,7 +4,6 @@ It is used by the "pytensor-cache" CLI tool, located in the /bin folder of the r
 """
 
 import logging
-import os
 import pickle
 import shutil
 from collections import Counter
@@ -33,12 +32,11 @@ def cleanup():
     If there is no key left for a compiled module, we delete the module.
 
     """
-    compiledir = config.compiledir
-    for directory in os.listdir(compiledir):
+    for directory in config.compiledir.iterdir():
         try:
-            filename = os.path.join(compiledir, directory, "key.pkl")
+            filename = directory / "key.pkl"
             # print file
-            with open(filename, "rb") as file:
+            with filename.open("rb") as file:
                 try:
                     keydata = pickle.load(file)
 
@@ -79,7 +77,7 @@ def cleanup():
                                     "the directory containing it."
                                 )
                     if len(keydata.keys) == 0:
-                        shutil.rmtree(os.path.join(compiledir, directory))
+                        shutil.rmtree(directory)
 
                 except (EOFError, AttributeError):
                     _logger.error(
@@ -117,11 +115,11 @@ def print_compiledir_content():
     big_key_files = []
     total_key_sizes = 0
     nb_keys = Counter()
-    for dir in os.listdir(compiledir):
-        filename = os.path.join(compiledir, dir, "key.pkl")
-        if not os.path.exists(filename):
+    for dir in config.compiledir.iterdir():
+        filename = dir / "key.pkl"
+        if not filename.exists():
             continue
-        with open(filename, "rb") as file:
+        with filename.open("rb") as file:
             try:
                 keydata = pickle.load(file)
                 ops = list({x for x in flatten(keydata.keys) if isinstance(x, Op)})
@@ -134,15 +132,11 @@ def print_compiledir_content():
                         {x for x in flatten(keydata.keys) if isinstance(x, CType)}
                     )
                     compile_start = compile_end = float("nan")
-                    for fn in os.listdir(os.path.join(compiledir, dir)):
-                        if fn.startswith("mod.c"):
-                            compile_start = os.path.getmtime(
-                                os.path.join(compiledir, dir, fn)
-                            )
-                        elif fn.endswith(".so"):
-                            compile_end = os.path.getmtime(
-                                os.path.join(compiledir, dir, fn)
-                            )
+                    for fn in dir.iterdir():
+                        if fn.name == "mod.c":
+                            compile_start = fn.stat().st_mtime
+                        elif fn.suffix == ".so":
+                            compile_end = fn.stat().st_mtime
                     compile_time = compile_end - compile_start
                     if len(ops) == 1:
                         table.append((dir, ops[0], types, compile_time))
@@ -153,7 +147,7 @@ def print_compiledir_content():
                             (dir, ops_to_str, types_to_str, compile_time)
                         )
 
-                size = os.path.getsize(filename)
+                size = filename.stat().st_size
                 total_key_sizes += size
                 if size > max_key_file_size:
                     big_key_files.append((dir, size, ops))
@@ -239,8 +233,8 @@ def basecompiledir_ls():
     """
     subdirs = []
     others = []
-    for f in os.listdir(config.base_compiledir):
-        if os.path.isdir(os.path.join(config.base_compiledir, f)):
+    for f in config.base_compiledir.iterdir():
+        if f.is_dir():
             subdirs.append(f)
         else:
             others.append(f)

{pytensor-2.24.2 → pytensor-2.25.0}/pytensor/compile/compilelock.py

@@ -6,6 +6,7 @@ in the same compilation directory (which can cause crashes).
 import os
 import threading
 from contextlib import contextmanager
+from pathlib import Path
 
 import filelock
 
@@ -35,7 +36,7 @@ def force_unlock(lock_dir: os.PathLike):
         Path to a directory that was locked with `lock_ctx`.
     """
 
-    fl = filelock.FileLock(os.path.join(lock_dir, ".lock"))
+    fl = filelock.FileLock(Path(lock_dir) / ".lock")
     fl.release(force=True)
 
     dir_key = f"{lock_dir}-{os.getpid()}"
@@ -72,7 +73,7 @@ def lock_ctx(
 
     if dir_key not in local_mem._locks:
         local_mem._locks[dir_key] = True
-        fl = filelock.FileLock(os.path.join(lock_dir, ".lock"))
+        fl = filelock.FileLock(Path(lock_dir) / ".lock")
         fl.acquire(timeout=timeout)
         try:
             yield

{pytensor-2.24.2 → pytensor-2.25.0}/pytensor/compile/debugmode.py

@@ -30,6 +30,7 @@ from pytensor.configdefaults import config
 from pytensor.graph.basic import Variable, io_toposort
 from pytensor.graph.destroyhandler import DestroyHandler
 from pytensor.graph.features import AlreadyThere, BadOptimization
+from pytensor.graph.fg import Output
 from pytensor.graph.op import HasInnerGraph, Op
 from pytensor.graph.utils import InconsistencyError, MethodNotDefined
 from pytensor.link.basic import Container, LocalLinker
@@ -628,7 +629,9 @@ def _is_used_in_graph(fgraph, var):
         True if `var` is used by another node in the graph.
 
     """
-    return not (fgraph.clients[var] == [("output", 1)] or fgraph.clients[var] == [])
+    return any(
+        client for client, _ in fgraph.clients[var] if not isinstance(client.op, Output)
+    )
 
 
 def _check_strides_match(a, b, warn_err, op):
@@ -977,7 +980,7 @@ def _check_preallocated_output(
     # disable memory checks in that mode, since they were already run.
     try:
         changed_inner_mode = False
-        if isinstance(getattr(node, "op", None), HasInnerGraph):
+        if isinstance(node.op, HasInnerGraph):
             fn = node.op.fn
             if not (fn and hasattr(fn, "maker") and hasattr(fn.maker, "mode")):
                 _logger.warning(f"Expected pytensor function not found in {node.op}.fn")
@@ -1132,18 +1135,14 @@ class _FunctionGraphEvent:
 
     def __init__(self, kind, node, idx=None, reason=None):
         self.kind = kind
-        if node == "output":
-            self.node = "output"
-            self.op = "output"
-        else:
-            self.node = node
-            self.op = node.op
+        self.node = node
+        self.op = node.op
         self.idx = idx
         self.reason = str(reason)
 
     def __str__(self):
         if self.kind == "change":
-            if self.op != "output":
+            if not isinstance(self.op, Output):
                 msg = str(len(self.node.inputs))
             else:
                 msg = ""

{pytensor-2.24.2 → pytensor-2.25.0}/pytensor/compile/function/__init__.py

@@ -1,9 +1,15 @@
 import logging
 import re
 import traceback as tb
+from collections.abc import Iterable
+from pathlib import Path
 
+import pytensor.misc.pkl_utils
 from pytensor.compile.function.pfunc import pfunc
 from pytensor.compile.function.types import orig_function
+from pytensor.compile.mode import Mode
+from pytensor.compile.profiling import ProfileStats
+from pytensor.graph import Variable
 
 
 __all__ = ["types", "pfunc"]
@@ -13,20 +19,24 @@ _logger = logging.getLogger("pytensor.compile.function")
 
 
 def function_dump(
-    filename,
-    inputs,
-    outputs=None,
-    mode=None,
-    updates=None,
-    givens=None,
-    no_default_updates=False,
-    accept_inplace=False,
-    name=None,
-    rebuild_strict=True,
-    allow_input_downcast=None,
-    profile=None,
-    on_unused_input=None,
-    extra_tag_to_remove=None,
+    filename: str | Path,
+    inputs: Iterable[Variable],
+    outputs: Variable | Iterable[Variable] | dict[str, Variable] | None = None,
+    mode: str | Mode | None = None,
+    updates: Iterable[tuple[Variable, Variable]]
+    | dict[Variable, Variable]
+    | None = None,
+    givens: Iterable[tuple[Variable, Variable]]
+    | dict[Variable, Variable]
+    | None = None,
+    no_default_updates: bool = False,
+    accept_inplace: bool = False,
+    name: str | None = None,
+    rebuild_strict: bool = True,
+    allow_input_downcast: bool | None = None,
+    profile: bool | ProfileStats | None = None,
+    on_unused_input: str | None = None,
+    extra_tag_to_remove: str | None = None,
 ):
     """
     This is helpful to make a reproducible case for problems during PyTensor
@@ -59,24 +69,21 @@ def function_dump(
     `['annotations', 'replacement_of', 'aggregation_scheme', 'roles']`
 
     """
-    assert isinstance(filename, str)
-    d = dict(
-        inputs=inputs,
-        outputs=outputs,
-        mode=mode,
-        updates=updates,
-        givens=givens,
-        no_default_updates=no_default_updates,
-        accept_inplace=accept_inplace,
-        name=name,
-        rebuild_strict=rebuild_strict,
-        allow_input_downcast=allow_input_downcast,
-        profile=profile,
-        on_unused_input=on_unused_input,
-    )
-    with open(filename, "wb") as f:
-        import pytensor.misc.pkl_utils
-
+    d = {
+        "inputs": inputs,
+        "outputs": outputs,
+        "mode": mode,
+        "updates": updates,
+        "givens": givens,
+        "no_default_updates": no_default_updates,
+        "accept_inplace": accept_inplace,
+        "name": name,
+        "rebuild_strict": rebuild_strict,
+        "allow_input_downcast": allow_input_downcast,
+        "profile": profile,
+        "on_unused_input": on_unused_input,
+    }
+    with Path(filename).open("wb") as f:
         pickler = pytensor.misc.pkl_utils.StripPickler(
             f, protocol=-1, extra_tag_to_remove=extra_tag_to_remove
         )
@@ -84,18 +91,22 @@ def function_dump(
 
 
 def function(
-    inputs,
-    outputs=None,
-    mode=None,
-    updates=None,
-    givens=None,
-    no_default_updates=False,
-    accept_inplace=False,
-    name=None,
-    rebuild_strict=True,
-    allow_input_downcast=None,
-    profile=None,
-    on_unused_input=None,
+    inputs: Iterable[Variable],
+    outputs: Variable | Iterable[Variable] | dict[str, Variable] | None = None,
+    mode: str | Mode | None = None,
+    updates: Iterable[tuple[Variable, Variable]]
+    | dict[Variable, Variable]
+    | None = None,
+    givens: Iterable[tuple[Variable, Variable]]
+    | dict[Variable, Variable]
+    | None = None,
+    no_default_updates: bool = False,
+    accept_inplace: bool = False,
+    name: str | None = None,
+    rebuild_strict: bool = True,
+    allow_input_downcast: bool | None = None,
+    profile: bool | ProfileStats | None = None,
+    on_unused_input: str | None = None,
 ):
     """
     Return a :class:`callable object <pytensor.compile.function.types.Function>`

{pytensor-2.24.2 → pytensor-2.25.0}/pytensor/compile/function/types.py

@@ -78,8 +78,6 @@ def view_tree_set(fgraph, v, treeset):
     """
     treeset.add(v)
     for cl, v_input_pos_to_cl in fgraph.clients[v]:
-        if cl == "output":
-            continue
         vmap = cl.op.view_map
         dmap = cl.op.destroy_map
         for opos, iposlist in chain(vmap.items(), dmap.items()):
@@ -661,7 +659,7 @@ class Function:
             exist_svs = [i.variable for i in maker.inputs]
 
             # Check if given ShareVariables exist
-            for sv in swap.keys():
+            for sv in swap:
                 if sv not in exist_svs:
                     raise ValueError(f"SharedVariable: {sv.name} not found")
 
@@ -713,9 +711,9 @@ class Function:
         # it is well tested, we don't share the part of the storage_map.
         if share_memory:
             i_o_vars = maker.fgraph.inputs + maker.fgraph.outputs
-            for key in storage_map.keys():
+            for key, val in storage_map.items():
                 if key not in i_o_vars:
-                    new_storage_map[memo[key]] = storage_map[key]
+                    new_storage_map[memo[key]] = val
 
         if not name and self.name:
             name = self.name + " copy"
@@ -1202,8 +1200,11 @@ def insert_deepcopy(fgraph, wrapped_inputs, wrapped_outputs):
     has_destroyers_attr = hasattr(fgraph, "has_destroyers")
 
     for i in range(len(fgraph.outputs)):
+        original_out = fgraph.outputs[i]
+        output_client = fgraph.get_output_client(i)
+
         views_of_output_i = set()
-        view_tree_set(fgraph, alias_root(fgraph.outputs[i]), views_of_output_i)
+        view_tree_set(fgraph, alias_root(original_out), views_of_output_i)
         copied = False
         # do not allow outputs to be aliased
         for j in range(i + 1, len(fgraph.outputs)):
@@ -1212,16 +1213,16 @@ def insert_deepcopy(fgraph, wrapped_inputs, wrapped_outputs):
             if fgraph.outputs[j] in views_of_output_i:
                 if wrapped_outputs[i].borrow and wrapped_outputs[j].borrow:
                     fgraph.change_node_input(
-                        "output", i, view_op(fgraph.outputs[i]), reason=reason
+                        *output_client, view_op(original_out), reason=reason
                     )
                 else:
                     fgraph.change_node_input(
-                        "output", i, deep_copy_op(fgraph.outputs[i]), reason=reason
+                        *output_client, deep_copy_op(original_out), reason=reason
                     )
                 copied = True
                 break
 
-        if not copied:
+        if not copied:  # no-break
             for input_j in all_graph_inputs:
                 # do not allow outputs to be aliased to an inputs (j), unless
                 # a) that j'th input has been 'destroyed' by
@@ -1239,33 +1240,29 @@ def insert_deepcopy(fgraph, wrapped_inputs, wrapped_outputs):
                         j = fgraph.inputs.index(input_j)
                         if wrapped_outputs[i].borrow and wrapped_inputs[j].borrow:
                             fgraph.change_node_input(
-                                "output",
-                                i,
-                                view_op(fgraph.outputs[i]),
+                                *output_client,
+                                view_op(original_out),
                                 reason=reason,
                             )
                             break
                         else:
                             fgraph.change_node_input(
-                                "output",
-                                i,
-                                deep_copy_op(fgraph.outputs[i]),
+                                *output_client,
+                                deep_copy_op(original_out),
                                 reason=reason,
                             )
                             break
                     elif wrapped_outputs[i].borrow:
                         fgraph.change_node_input(
-                            "output",
-                            i,
-                            view_op(fgraph.outputs[i]),
+                            *output_client,
+                            view_op(original_out),
                             reason=reason,
                         )
                         break
                     else:
                         fgraph.change_node_input(
-                            "output",
-                            i,
-                            deep_copy_op(fgraph.outputs[i]),
+                            *output_client,
+                            deep_copy_op(original_out),
                             reason=reason,
                         )
                         break
@@ -1449,7 +1446,7 @@ class FunctionMaker:
         if not hasattr(mode.linker, "accept"):
             raise ValueError(
                 "'linker' parameter of FunctionMaker should be "
-                f"a Linker with an accept method or one of {list(pytensor.compile.mode.predefined_linkers.keys())}"
+                f"a Linker with an accept method or one of {list(pytensor.compile.mode.predefined_linkers)}"
             )
 
     def __init__(