pytensor 3.3.0__tar.gz → 3.3.1__tar.gz

{pytensor-3.3.0/pytensor.egg-info → pytensor-3.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.3.0
+Version: 3.3.1
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause
@@ -31,7 +31,7 @@ License-File: LICENSE.txt
 Requires-Dist: setuptools>=59.0.0
 Requires-Dist: scipy<2,>=1
 Requires-Dist: numpy>=2.0
-Requires-Dist: numba<=0.66.0,>=0.58
+Requires-Dist: numba<=0.67.0,>=0.58
 Requires-Dist: filelock>=3.15
 Provides-Extra: complete
 Requires-Dist: pytensor[jax]; extra == "complete"
@@ -57,7 +57,7 @@ Provides-Extra: jax
 Requires-Dist: jax; extra == "jax"
 Requires-Dist: jaxlib; extra == "jax"
 Provides-Extra: numba
-Requires-Dist: numba<=0.66.0,>=0.58; extra == "numba"
+Requires-Dist: numba<=0.67.0,>=0.58; extra == "numba"
 Requires-Dist: llvmlite; extra == "numba"
 Provides-Extra: kanren
 Requires-Dist: etuples; extra == "kanren"

{pytensor-3.3.0 → pytensor-3.3.1}/doc/library/graph/index.rst

@@ -17,5 +17,6 @@
     replace
     features
     op
+    traversal
     type
     utils

pytensor-3.3.1/doc/library/graph/traversal.rst

@@ -0,0 +1,8 @@
+.. _libdoc_graph_traversal:
+
+================================================
+:mod:`traversal` -- Graph Traversal Utilities
+================================================
+
+.. automodule:: pytensor.graph.traversal
+   :members:

{pytensor-3.3.0 → pytensor-3.3.1}/pyproject.toml

@@ -50,7 +50,7 @@ dependencies = [
     "setuptools>=59.0.0",
     "scipy>=1,<2",
     "numpy>=2.0",
-    "numba>=0.58,<=0.66.0",
+    "numba>=0.58,<=0.67.0",
     "filelock>=3.15",
 ]
 
@@ -78,7 +78,7 @@ tests = [
 ]
 rtd = ["sphinx>=5.1.0,<6", "pygments", "pydot"]
 jax = ["jax", "jaxlib"]
-numba = ["numba>=0.58,<=0.66.0", "llvmlite"]
+numba = ["numba>=0.58,<=0.67.0", "llvmlite"]
 kanren = [
     "etuples",
     "logical-unification",
@@ -131,7 +131,11 @@ filterwarnings =[
 
 [tool.ruff]
 line-length = 88
-exclude = ["pytensor/_version.py"]
+# _patch_list_to_tuple.py is copied verbatim from numba; keep it diffable upstream
+exclude = [
+    "pytensor/_version.py",
+    "pytensor/link/numba/dispatch/_patch_list_to_tuple.py",
+]
 
 [tool.ruff.format]
 docstring-code-format = true

{pytensor-3.3.0 → pytensor-3.3.1}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-08-12T10:16:27+0200",
+ "date": "2026-09-07T16:04:36+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "aa1bc7772eb11ed8def2814374aeb5da03da590c",
- "version": "3.3.0"
+ "full-revisionid": "bb5967ab72cfc29ee9f47c4a2fd0aaeeb1695268",
+ "version": "3.3.1"
 }
 '''  # END VERSION_JSON
 

{pytensor-3.3.0 → pytensor-3.3.1}/pytensor/compile/mode.py

@@ -524,7 +524,9 @@ JAX = Mode(
 )
 PYTORCH = Mode(
     PytorchLinker(),
-    RewriteDatabaseQuery(include=["fast_run"]),
+    # `Composite` has no PyTorch dispatch (it has no `nfunc_spec` to look up in torch),
+    # and torch fuses its own kernels anyway.
+    RewriteDatabaseQuery(include=["fast_run"], exclude=["fusion"]),
 )
 
 MLX = Mode(

{pytensor-3.3.0 → pytensor-3.3.1}/pytensor/gradient.py

@@ -1335,8 +1335,8 @@ def _populate_grad_dict(var_to_app_to_idx, grad_dict, wrt, cost_name=None):
                     raise ValueError(
                         f"{node.op} returned the wrong number of gradient terms."
                     )
-            # We can not enforce this, as AdvancedSubtensor1 has an option to
-            # return the sparse grad for optimization reason.
+            # We can not enforce this, as sparse Ops have an option to return a
+            # structured grad for optimization reason.
 
             #            for ig, i in zip(input_grads, inputs):
             #                if (not isinstance(ig.type, (DisconnectedType, NullType)) and

{pytensor-3.3.0 → pytensor-3.3.1}/pytensor/graph/basic.py

@@ -45,13 +45,12 @@ _MOVED_FUNCTIONS = {
     "graph_inputs",
     "explicit_graph_inputs",
     "vars_between",
-    "orhpans_between",
+    "orphans_between",
     "applys_between",
     "apply_depends_on",
     "truncated_graph_inputs",
     "general_toposort",
     "io_toposort",
-    "list_of_nodes",
     "get_var_by_name",
 }
 

{pytensor-3.3.0 → pytensor-3.3.1}/pytensor/graph/replace.py

@@ -123,8 +123,6 @@ def graph_replace(
         Replace mapping
     strict: bool
         Raise an error if some replacements were not used
-    return_unused: bool
-        Return replacements that were not used
 
     Returns
     -------

{pytensor-3.3.0 → pytensor-3.3.1}/pytensor/graph/rewriting/basic.py

@@ -43,7 +43,6 @@ from pytensor.graph.traversal import (
 )
 from pytensor.graph.utils import AssocList, InconsistencyError
 from pytensor.misc.ordered_set import OrderedSet
-from pytensor.utils import flatten
 
 
 _logger = logging.getLogger("pytensor.graph.rewriting.basic")
@@ -705,6 +704,15 @@ class MergeOptimizer(GraphRewriter):
             fgraph.attach_feature(MergeFeature())
 
     def apply(self, fgraph):
+        def has_direct_destroyer(var):
+            """Whether any direct client destroys `var` (view-chain destroyers not seen)."""
+            return any(
+                i in destroyed_idxs
+                for client, i in fgraph.clients[var]
+                if client.op.destroy_map
+                for destroyed_idxs in client.op.destroy_map.values()
+            )
+
         sched = fgraph.merge_feature.scheduled
         nb_fail = 0
         t0 = time.perf_counter()
@@ -755,15 +763,17 @@ class MergeOptimizer(GraphRewriter):
                         continue
 
                     if hasattr(fgraph, "destroy_handler"):
-                        # If both nodes have clients that destroy them, we
-                        # can't merge them.
-                        clients = (
-                            fgraph.clients[pairs[0][0]] + fgraph.clients[pairs[0][1]]
-                        )
-                        if any(
-                            i in flatten(c.op.destroy_map.values())
-                            for c, i in clients
-                            if c.op.destroy_map
+                        # Merging two variables that both have a destroyer is
+                        # always rejected by `DestroyHandler` validation
+                        # ("multiple destroyers") -- skip those outright
+                        # instead of paying a checkpoint + failed validate +
+                        # revert per scheduled pair. A single-destroyer merge
+                        # may still be feasible, and destroyers reachable only
+                        # through view chains are rare: both are left to
+                        # validation, keeping this check O(clients) with no
+                        # droot recomputation.
+                        if has_direct_destroyer(pairs[0][0]) and has_direct_destroyer(
+                            pairs[0][1]
                         ):
                             continue
 

{pytensor-3.3.0 → pytensor-3.3.1}/pytensor/graph/traversal.py

@@ -145,12 +145,13 @@ def apply_ancestors(
     graphs: Iterable[Variable],
     blockers: Iterable[Variable] | None = None,
 ) -> Generator[Apply, None, None]:
-    """Return the Apply nodes that contribute to those in given graphs (inclusive)."""
+    """Return the Apply nodes that contribute to those in given graphs, excluding the owners of the blockers."""
     seen = {None}  # This filters out Variables without an owner
+    blockers = frozenset(blockers or ())
     for var in ancestors(graphs, blockers):
         # For multi-output nodes, we'll see multiple variables
         # but we should only yield the Apply once
-        if (apply := var.owner) not in seen:
+        if var not in blockers and (apply := var.owner) not in seen:
             yield apply
             seen.add(apply)
     return

{pytensor-3.3.0 → pytensor-3.3.1}/pytensor/link/mlx/dispatch/linalg/solvers.py

@@ -4,6 +4,7 @@ import mlx.core as mx
 
 from pytensor.link.mlx.dispatch.basic import mlx_funcify
 from pytensor.tensor.linalg.solvers.general import Solve
+from pytensor.tensor.linalg.solvers.psd import CholeskySolve
 from pytensor.tensor.linalg.solvers.triangular import SolveTriangular
 
 
@@ -44,3 +45,22 @@ def mlx_funcify_SolveTriangular(op, node, **kwargs):
         )
 
     return solve_triangular
+
+
+@mlx_funcify.register(CholeskySolve)
+def mlx_funcify_CholeskySolve(op, node, **kwargs):
+    lower = op.lower
+    c_dtype = getattr(mx, node.inputs[0].dtype)
+    b_dtype = getattr(mx, node.inputs[1].dtype)
+
+    # MLX has no cho_solve, so with A = L L.T we solve L y = b then L.T x = y.
+    def cho_solve(c, b):
+        c = c.astype(stream=mx.cpu, dtype=c_dtype)
+        b = b.astype(stream=mx.cpu, dtype=b_dtype)
+        c_T = mx.swapaxes(c, -1, -2, stream=mx.cpu)
+        L, L_T = (c, c_T) if lower else (c_T, c)
+
+        y = mx.linalg.solve_triangular(L, b, upper=False, stream=mx.cpu)
+        return mx.linalg.solve_triangular(L_T, y, upper=True, stream=mx.cpu)
+
+    return cho_solve

pytensor-3.3.1/pytensor/link/mlx/dispatch/scalar/__init__.py

@@ -0,0 +1,10 @@
+from pytensor.link.mlx.dispatch.scalar import (
+    basic,
+    bessel,
+    beta,
+    erf,
+    erfcinv,
+    gamma,
+    gammainc,
+    math,
+)