pytensor 2.23.0__tar.gz → 2.24.0__tar.gz

{pytensor-2.23.0 → pytensor-2.24.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytensor
-Version: 2.23.0
+Version: 2.24.0
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License: .. _license:
@@ -72,15 +72,14 @@ Classifier: Programming Language :: Python :: 3.12
 Requires-Python: <3.13,>=3.10
 Description-Content-Type: text/x-rst
 License-File: LICENSE.txt
-Requires-Dist: setuptools>=48.0.0
-Requires-Dist: scipy>=0.14
+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: etuples
 Requires-Dist: logical-unification
 Requires-Dist: miniKanren
 Requires-Dist: cons
-Requires-Dist: typing_extensions
 Provides-Extra: complete
 Requires-Dist: pytensor[jax]; extra == "complete"
 Requires-Dist: pytensor[numba]; extra == "complete"
@@ -106,6 +105,7 @@ Requires-Dist: jax; extra == "jax"
 Requires-Dist: jaxlib; extra == "jax"
 Provides-Extra: numba
 Requires-Dist: numba>=0.57; extra == "numba"
+Requires-Dist: llvmlite; extra == "numba"
 
 .. image:: https://cdn.rawgit.com/pymc-devs/pytensor/main/doc/images/PyTensor_RGB.svg
     :height: 100px

{pytensor-2.23.0 → pytensor-2.24.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [build-system]
 requires = [
-    "setuptools>=48.0.0",
+    "setuptools>=59.0.0",
     "cython",
     "numpy>=1.17.0",
     "versioneer[toml]>=0.28",
@@ -46,15 +46,14 @@ keywords = [
     "differentiation",
 ]
 dependencies = [
-    "setuptools>=48.0.0",
-    "scipy>=0.14",
+    "setuptools>=59.0.0",
+    "scipy>=0.14,<1.14",
     "numpy>=1.17.0,<2",
     "filelock",
     "etuples",
     "logical-unification",
     "miniKanren",
     "cons",
-    "typing_extensions",
 ]
 
 [project.urls]
@@ -79,7 +78,7 @@ tests = [
 ]
 rtd = ["sphinx>=5.1.0,<6", "pygments", "pydot", "pydot2", "pydot-ng"]
 jax = ["jax", "jaxlib"]
-numba = ["numba>=0.57"]
+numba = ["numba>=0.57", "llvmlite"]
 
 [tool.setuptools.packages.find]
 include = ["pytensor*"]

{pytensor-2.23.0 → pytensor-2.24.0}/pytensor/__init__.py

@@ -108,7 +108,7 @@ def as_symbolic(x: Any, name: str | None = None, **kwargs) -> Variable:
 
 
 @singledispatch
-def _as_symbolic(x, **kwargs) -> Variable:
+def _as_symbolic(x: Any, **kwargs) -> Variable:
     from pytensor.tensor import as_tensor_variable
 
     return as_tensor_variable(x, **kwargs)

{pytensor-2.23.0 → pytensor-2.24.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-06-06T14:05:58+0200",
+ "date": "2024-07-03T14:50:56+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "4b6a444003e3e2de82655818c5935f143a27a544",
- "version": "2.23.0"
+ "full-revisionid": "df769f6c75d8598de4be08f8053831e39263d96b",
+ "version": "2.24.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.23.0 → pytensor-2.24.0}/pytensor/compile/compiledir.py

@@ -7,6 +7,7 @@ import logging
 import os
 import pickle
 import shutil
+from collections import Counter
 
 import numpy as np
 
@@ -111,11 +112,11 @@ def print_compiledir_content():
     compiledir = config.compiledir
     table = []
     table_multiple_ops = []
-    table_op_class = {}
+    table_op_class = Counter()
     zeros_op = 0
     big_key_files = []
     total_key_sizes = 0
-    nb_keys = {}
+    nb_keys = Counter()
     for dir in os.listdir(compiledir):
         filename = os.path.join(compiledir, dir, "key.pkl")
         if not os.path.exists(filename):
@@ -125,9 +126,7 @@ def print_compiledir_content():
                 keydata = pickle.load(file)
                 ops = list({x for x in flatten(keydata.keys) if isinstance(x, Op)})
                 # Whatever the case, we count compilations for OP classes.
-                for op_class in {op.__class__ for op in ops}:
-                    table_op_class.setdefault(op_class, 0)
-                    table_op_class[op_class] += 1
+                table_op_class.update({op.__class__ for op in ops})
                 if len(ops) == 0:
                     zeros_op += 1
                 else:
@@ -159,7 +158,6 @@ def print_compiledir_content():
                 if size > max_key_file_size:
                     big_key_files.append((dir, size, ops))
 
-                nb_keys.setdefault(len(keydata.keys), 0)
                 nb_keys[len(keydata.keys)] += 1
             except OSError:
                 pass
@@ -198,8 +196,7 @@ def print_compiledir_content():
         ),
         underline="+",
     )
-    table_op_class = sorted(table_op_class.items(), key=lambda t: t[1])
-    for op_class, nb in table_op_class:
+    for op_class, nb in reversed(table_op_class.most_common()):
         print(op_class, nb)
 
     if big_key_files:

{pytensor-2.23.0 → pytensor-2.24.0}/pytensor/compile/debugmode.py

@@ -687,7 +687,7 @@ def _lessbroken_deepcopy(a):
     else:
         rval = copy.deepcopy(a)
 
-    assert type(rval) == type(a), (type(rval), type(a))
+    assert type(rval) is type(a), (type(rval), type(a))
 
     if isinstance(rval, np.ndarray):
         assert rval.dtype == a.dtype
@@ -906,11 +906,10 @@ def _get_preallocated_maps(
                 name = f"strided{tuple(steps)}"
                 for r in considered_outputs:
                     if r in init_strided:
-                        strides = []
-                        shapes = []
-                        for i, size in enumerate(r_vals[r].shape):
-                            shapes.append(slice(None, size, None))
-                            strides.append(slice(None, None, steps[i]))
+                        shapes = [slice(None, size, None) for size in r_vals[r].shape]
+                        strides = [
+                            slice(None, None, steps[i]) for i in range(r_vals[r].ndim)
+                        ]
 
                         r_buf = init_strided[r]
 
@@ -1154,7 +1153,7 @@ class _FunctionGraphEvent:
             return str(self.__dict__)
 
     def __eq__(self, other):
-        rval = type(self) == type(other)
+        rval = type(self) is type(other)
         if rval:
             # nodes are not compared because this comparison is
             # supposed to be true for corresponding events that happen

{pytensor-2.23.0 → pytensor-2.24.0}/pytensor/compile/function/__init__.py

@@ -247,18 +247,10 @@ def function(
 
     """
     if isinstance(outputs, dict):
-        output_items = list(outputs.items())
+        assert all(isinstance(k, str) for k in outputs)
 
-        for item_pair in output_items:
-            assert isinstance(item_pair[0], str)
-
-        output_items_sorted = sorted(output_items)
-
-        output_keys = []
-        outputs = []
-        for pair in output_items_sorted:
-            output_keys.append(pair[0])
-            outputs.append(pair[1])
+        output_keys = sorted(outputs)
+        outputs = [outputs[key] for key in output_keys]
 
     else:
         output_keys = None

{pytensor-2.23.0 → pytensor-2.24.0}/pytensor/compile/function/types.py

@@ -14,6 +14,7 @@ import pytensor
 import pytensor.compile.profiling
 from pytensor.compile.io import In, SymbolicInput, SymbolicOutput
 from pytensor.compile.ops import deep_copy_op, view_op
+from pytensor.compile.profiling import ProfileStats
 from pytensor.configdefaults import config
 from pytensor.graph.basic import (
     Constant,
@@ -212,18 +213,14 @@ def std_fgraph(
 
         found_updates.extend(map(SymbolicOutput, updates))
     elif fgraph is None:
-        input_vars = []
-
         # If one of the inputs is non-atomic (i.e. has a non-`None` `Variable.owner`),
         # then we need to create/clone the graph starting at these inputs.
         # The result will be atomic versions of the given inputs connected to
         # the same outputs.
         # Otherwise, when all the inputs are already atomic, there's no need to
         # clone the graph.
-        clone = force_clone
-        for spec in input_specs:
-            input_vars.append(spec.variable)
-            clone |= spec.variable.owner is not None
+        input_vars = [spec.variable for spec in input_specs]
+        clone = force_clone or any(var.owner is not None for var in input_vars)
 
         fgraph = FunctionGraph(
             input_vars,
@@ -557,11 +554,11 @@ class Function:
 
     def copy(
         self,
-        share_memory=False,
-        swap=None,
-        delete_updates=False,
-        name=None,
-        profile=None,
+        share_memory: bool = False,
+        swap: dict | None = None,
+        delete_updates: bool = False,
+        name: str | None = None,
+        profile: bool | str | ProfileStats | None = None,
     ):
         """
         Copy this function. Copied function will have separated maker and
@@ -588,7 +585,7 @@ class Function:
             If provided, will be the name of the new
             Function. Otherwise, it will be old + " copy"
 
-        profile :
+        profile : bool | str | ProfileStats | None
             as pytensor.function profile parameter
 
         Returns
@@ -727,14 +724,8 @@ class Function:
         # reinitialize new maker and create new function
         if profile is None:
             profile = config.profile or config.print_global_stats
-            # profile -> True or False
         if profile is True:
-            if name:
-                message = name
-            else:
-                message = str(profile.message) + " copy"
-            profile = pytensor.compile.profiling.ProfileStats(message=message)
-            # profile -> object
+            profile = pytensor.compile.profiling.ProfileStats(message=name)
         elif isinstance(profile, str):
             profile = pytensor.compile.profiling.ProfileStats(message=profile)
 
@@ -1106,6 +1097,18 @@ class Function:
         # NOTE: sync was needed on old gpu backend
         pass
 
+    def dprint(self, **kwargs):
+        """Debug print itself
+
+        Parameters
+        ----------
+        kwargs:
+            Optional keyword arguments to pass to debugprint function.
+        """
+        from pytensor.printing import debugprint
+
+        return debugprint(self, **kwargs)
+
 
 # pickling/deepcopy support for Function
 def _pickle_Function(f):

{pytensor-2.23.0 → pytensor-2.24.0}/pytensor/compile/mode.py

@@ -28,6 +28,7 @@ from pytensor.link.basic import Linker, PerformLinker
 from pytensor.link.c.basic import CLinker, OpWiseCLinker
 from pytensor.link.jax.linker import JAXLinker
 from pytensor.link.numba.linker import NumbaLinker
+from pytensor.link.pytorch.linker import PytorchLinker
 from pytensor.link.vm import VMLinker
 
 
@@ -47,6 +48,7 @@ predefined_linkers = {
     "vm_nogc": VMLinker(allow_gc=False, use_cloop=False),
     "cvm_nogc": VMLinker(allow_gc=False, use_cloop=True),
     "jax": JAXLinker(),
+    "pytorch": PytorchLinker(),
     "numba": NumbaLinker(),
 }
 
@@ -460,6 +462,18 @@ JAX = Mode(
         ],
     ),
 )
+PYTORCH = Mode(
+    PytorchLinker(),
+    RewriteDatabaseQuery(
+        include=["fast_run"],
+        exclude=[
+            "cxx_only",
+            "BlasOpt",
+            "fusion",
+            "inplace",
+        ],
+    ),
+)
 NUMBA = Mode(
     NumbaLinker(),
     RewriteDatabaseQuery(
@@ -474,6 +488,7 @@ predefined_modes = {
     "FAST_RUN": FAST_RUN,
     "JAX": JAX,
     "NUMBA": NUMBA,
+    "PYTORCH": PYTORCH,
 }
 
 instantiated_default_mode = None

{pytensor-2.23.0 → pytensor-2.24.0}/pytensor/compile/ops.py

@@ -246,7 +246,7 @@ class FromFunctionOp(Op):
             self.infer_shape = self._infer_shape
 
     def __eq__(self, other):
-        return type(self) == type(other) and self.__fn == other.__fn
+        return type(self) is type(other) and self.__fn == other.__fn
 
     def __hash__(self):
         return hash(type(self)) ^ hash(self.__fn)

{pytensor-2.23.0 → pytensor-2.24.0}/pytensor/compile/profiling.py

@@ -14,9 +14,9 @@ import logging
 import operator
 import sys
 import time
-from collections import defaultdict
+from collections import Counter, defaultdict
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
@@ -204,8 +204,8 @@ class ProfileStats:
         self.fct_call_time = 0.0
         self.fct_callcount = 0
         self.vm_call_time = 0.0
-        self.apply_time = {}
-        self.apply_callcount = {}
+        self.apply_time = defaultdict(float)
+        self.apply_callcount = Counter()
         # self.apply_cimpl = None
         # self.message = None
 
@@ -234,9 +234,9 @@ class ProfileStats:
     # Total time spent in Function.vm.__call__
     #
 
-    apply_time: dict[Union["FunctionGraph", Variable], float] | None = None
+    apply_time: dict[tuple["FunctionGraph", Apply], float]
 
-    apply_callcount: dict[Union["FunctionGraph", Variable], int] | None = None
+    apply_callcount: dict[tuple["FunctionGraph", Apply], int]
 
     apply_cimpl: dict[Apply, bool] | None = None
     # dict from node -> bool (1 if c, 0 if py)
@@ -292,10 +292,9 @@ class ProfileStats:
     # param is called flag_time_thunks because most other attributes with time
     # in the name are times *of* something, rather than configuration flags.
     def __init__(self, atexit_print=True, flag_time_thunks=None, **kwargs):
-        self.apply_callcount = {}
+        self.apply_callcount = Counter()
         self.output_size = {}
-        # Keys are `(FunctionGraph, Variable)`
-        self.apply_time = {}
+        self.apply_time = defaultdict(float)
         self.apply_cimpl = {}
         self.variable_shape = {}
         self.variable_strides = {}
@@ -320,12 +319,10 @@ class ProfileStats:
 
         """
         # timing is stored by node, we compute timing by class on demand
-        rval = {}
-        for (fgraph, node), t in self.apply_time.items():
-            typ = type(node.op)
-            rval.setdefault(typ, 0)
-            rval[typ] += t
-        return rval
+        rval = defaultdict(float)
+        for (_fgraph, node), t in self.apply_time.items():
+            rval[type(node.op)] += t
+        return dict(rval)
 
     def class_callcount(self):
         """
@@ -333,24 +330,18 @@ class ProfileStats:
 
         """
         # timing is stored by node, we compute timing by class on demand
-        rval = {}
-        for (fgraph, node), count in self.apply_callcount.items():
-            typ = type(node.op)
-            rval.setdefault(typ, 0)
-            rval[typ] += count
+        rval = Counter()
+        for (_fgraph, node), count in self.apply_callcount.items():
+            rval[type(node.op)] += count
         return rval
 
-    def class_nodes(self):
+    def class_nodes(self) -> Counter:
         """
         dict op -> total number of nodes
 
         """
         # timing is stored by node, we compute timing by class on demand
-        rval = {}
-        for (fgraph, node), count in self.apply_callcount.items():
-            typ = type(node.op)
-            rval.setdefault(typ, 0)
-            rval[typ] += 1
+        rval = Counter(type(node.op) for _fgraph, node in self.apply_callcount)
         return rval
 
     def class_impl(self):
@@ -360,12 +351,9 @@ class ProfileStats:
         """
         # timing is stored by node, we compute timing by class on demand
         rval = {}
-        for fgraph, node in self.apply_callcount:
+        for _fgraph, node in self.apply_callcount:
             typ = type(node.op)
-            if self.apply_cimpl[node]:
-                impl = "C "
-            else:
-                impl = "Py"
+            impl = "C " if self.apply_cimpl[node] else "Py"
             rval.setdefault(typ, impl)
             if rval[typ] != impl and len(rval[typ]) == 2:
                 rval[typ] += impl
@@ -377,11 +365,10 @@ class ProfileStats:
 
         """
         # timing is stored by node, we compute timing by Op on demand
-        rval = {}
+        rval = defaultdict(float)
         for (fgraph, node), t in self.apply_time.items():
-            rval.setdefault(node.op, 0)
             rval[node.op] += t
-        return rval
+        return dict(rval)
 
     def fill_node_total_time(self, fgraph, node, total_times):
         """
@@ -414,9 +401,8 @@ class ProfileStats:
 
         """
         # timing is stored by node, we compute timing by Op on demand
-        rval = {}
+        rval = Counter()
         for (fgraph, node), count in self.apply_callcount.items():
-            rval.setdefault(node.op, 0)
             rval[node.op] += count
         return rval
 
@@ -426,10 +412,7 @@ class ProfileStats:
 
         """
         # timing is stored by node, we compute timing by Op on demand
-        rval = {}
-        for (fgraph, node), count in self.apply_callcount.items():
-            rval.setdefault(node.op, 0)
-            rval[node.op] += 1
+        rval = Counter(node.op for _fgraph, node in self.apply_callcount)
         return rval
 
     def op_impl(self):
@@ -1204,8 +1187,7 @@ class ProfileStats:
                         compute_map[var][0] = 0
 
                     for k_remove, v_remove in viewedby_remove.items():
-                        for i in v_remove:
-                            viewed_by[k_remove].append(i)
+                        viewed_by[k_remove].extend(v_remove)
 
                     for k_add, v_add in viewedby_add.items():
                         for i in v_add:
@@ -1215,15 +1197,16 @@ class ProfileStats:
                         del view_of[k]
 
             # two data structure used to mimic Python gc
-            viewed_by = {}  # {var1: [vars that view var1]}
+            # * {var1: [vars that view var1]}
             # The len of the list is the value of python ref
             # count. But we use a list, not just the ref count value.
-            # This is more safe to help detect potential bug  in the algo
-            for var in fgraph.variables:
-                viewed_by[var] = []
-            view_of = {}  # {var1: original var viewed by var1}
+            # This is more safe to help detect potential bug in the algo
+            viewed_by = {var: [] for var in fgraph.variables}
+
+            # * {var1: original var viewed by var1}
             # The original mean that we don't keep track of all the intermediate
             # relationship in the view.
+            view_of = {}
 
             min_memory_generator(executable_nodes, viewed_by, view_of)
 

{pytensor-2.23.0 → pytensor-2.24.0}/pytensor/gradient.py

@@ -1041,13 +1041,12 @@ def _populate_grad_dict(var_to_app_to_idx, grad_dict, wrt, cost_name=None):
             # list of bools indicating if each input is connected to the cost
             inputs_connected = [
                 (
-                    True
-                    in [
+                    any(
                         input_to_output and output_to_cost
                         for input_to_output, output_to_cost in zip(
                             input_to_outputs, outputs_connected
                         )
-                    ]
+                    )
                 )
                 for input_to_outputs in connection_pattern
             ]
@@ -1067,25 +1066,24 @@ def _populate_grad_dict(var_to_app_to_idx, grad_dict, wrt, cost_name=None):
             # List of bools indicating if each input only has NullType outputs
             only_connected_to_nan = [
                 (
-                    True
-                    not in [
+                    not any(
                         in_to_out and out_to_cost and not out_nan
                         for in_to_out, out_to_cost, out_nan in zip(
                             in_to_outs, outputs_connected, ograd_is_nan
                         )
-                    ]
+                    )
                 )
                 for in_to_outs in connection_pattern
             ]
 
-            if True not in inputs_connected:
+            if not any(inputs_connected):
                 # All outputs of this op are disconnected so we can skip
                 # Calling the op's grad method and report that the inputs
                 # are disconnected
                 # (The op's grad method could do this too, but this saves the
                 # implementer the trouble of worrying about this case)
                 input_grads = [disconnected_type() for ipt in inputs]
-            elif False not in only_connected_to_nan:
+            elif all(only_connected_to_nan):
                 # All inputs are only connected to nan gradients, so we don't
                 # need to bother calling the grad method. We know the gradient
                 # with respect to all connected inputs is nan.

{pytensor-2.23.0 → pytensor-2.24.0}/pytensor/graph/basic.py

@@ -748,7 +748,7 @@ class NominalVariable(AtomicVariable[_TypeType]):
             return True
 
         return (
-            type(self) == type(other)
+            type(self) is type(other)
             and self.id == other.id
             and self.type == other.type
         )
@@ -1302,8 +1302,8 @@ def clone_node_and_cache(
 
 
 def clone_get_equiv(
-    inputs: Sequence[Variable],
-    outputs: Sequence[Variable],
+    inputs: Iterable[Variable],
+    outputs: Reversible[Variable],
     copy_inputs: bool = True,
     copy_orphans: bool = True,
     memo: dict[Union[Apply, Variable, "Op"], Union[Apply, Variable, "Op"]]
@@ -1474,9 +1474,8 @@ def general_toposort(
 
     _clients: dict[T, list[T]] = {}
     sources: deque[T] = deque()
-    search_res_len: int = 0
+    search_res_len = len(search_res)
     for snode, children in search_res:
-        search_res_len += 1
         if children:
             for child in children:
                 _clients.setdefault(child, []).append(snode)