pytensor 3.2.4__tar.gz → 3.3.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.2.4
+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
@@ -24,13 +24,14 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python :: Free Threading :: 1 - Unstable
 Requires-Python: <3.15,>=3.12
 Description-Content-Type: text/x-rst
 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"
@@ -56,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.2.4 → 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.2.4 → pytensor-3.3.1}/pyproject.toml

@@ -33,6 +33,7 @@ classifiers = [
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
     "Programming Language :: Python :: 3.14",
+    "Programming Language :: Python :: Free Threading :: 1 - Unstable",
 ]
 
 keywords = [
@@ -49,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",
 ]
 
@@ -77,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",
@@ -130,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
@@ -185,11 +190,18 @@ files = ["pytensor", "tests"]
 build = "*"
 # Uncomment to skip builds that compile but fail when trying to test (maybe due to incompatibility with runner)
 # archs = ["auto64"]
-# Disable any-platform (pp*), 32-bit builds, and free-threaded builds (cp*t-*)
+# Disable any-platform (pp*) and 32-bit builds
 # Additional options to consider: "*musllinux*"
-skip = ["pp*", "*-win32", "*-manylinux_i686", "cp*t-*"]
+skip = ["pp*", "*-win32", "*-manylinux_i686"]
+# Free-threaded (cp*t-*) wheels are built by default. The compiled extensions
+# are not marked free-threading-safe, so importing them re-enables the GIL
+# (with a warning). That's fine: they're no longer loaded by default.
 build-frontend = "build"
-test-command = 'python -c "import pytensor; print(pytensor.__version__); from pytensor.scan import scan_perform; print(scan_perform.get_version())"'
+# Import the compiled extensions and evaluate a function to check the wheel loads and
+# runs everywhere, including free-threaded builds. cxx="" disables the C backend (the
+# minimal wheel-test env has no libpython to link a runtime-compiled thunk against) and
+# FAST_COMPILE avoids numba, which crashes on free-threaded Windows.
+test-command = '''python -c "import pytensor; pytensor.config.cxx = ''; import pytensor.tensor as pt; print(pytensor.__version__); from pytensor.scan import scan_perform; print(scan_perform.get_version()); x = pt.vector(); f = pytensor.function([x], (x * 2).sum(), mode='FAST_COMPILE'); assert float(f([1.0, 2.0, 3.0])) == 12.0"'''
 test-skip = ["*musllinux*", "*i686*"]
 
 # Testing seems to be running into issues locating libs where expected

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

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-08-02T00:12:33+0200",
+ "date": "2026-09-07T16:04:36+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "7175f7cadf3934a03f63e6e7b187fa1519e6d85f",
- "version": "3.2.4"
+ "full-revisionid": "bb5967ab72cfc29ee9f47c4a2fd0aaeeb1695268",
+ "version": "3.3.1"
 }
 '''  # END VERSION_JSON
 

{pytensor-3.2.4 → 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.2.4 → 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.2.4 → 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.2.4 → 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.2.4 → 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.2.4 → 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.2.4 → pytensor-3.3.1}/pytensor/link/jax/dispatch/subtensor.py

@@ -1,3 +1,5 @@
+import jax.numpy as jnp
+
 from pytensor.link.jax.dispatch.basic import jax_funcify
 from pytensor.tensor.subtensor import (
     AdvancedIncSubtensor,
@@ -60,7 +62,11 @@ def jax_funcify_IncSubtensor(op, node, **kwargs):
             indices = indices[0]
 
         if isinstance(op, AdvancedIncSubtensor):
-            op._check_runtime_broadcast_of_vector_index(node, x, y, indices)
+            # jax_typify downgrades 0d arrays to Python scalars, which have no .shape,
+            # so re-arrayify y for a check that is written against array values.
+            op._check_runtime_broadcast_of_vector_index(
+                node, x, jnp.asarray(y), indices
+            )
 
         return jax_fn(x, indices, y)
 

{pytensor-3.2.4 → pytensor-3.3.1}/pytensor/link/mlx/dispatch/basic.py

@@ -15,25 +15,36 @@ from pytensor.link.utils import fgraph_to_python
 from pytensor.raise_op import Assert, CheckAndRaise
 
 
+def float64_supported():
+    """Return whether the current default device can operate on float64.
+
+    MLX implements float64 on the CPU only; the Metal backend rejects it outright.
+    """
+    return mx.default_device() == mx.cpu
+
+
 def convert_dtype_to_mlx(dtype_str, auto_cast_unsupported=True):
     """Convert PyTensor dtype strings to MLX dtype objects.
 
     MLX expects dtype objects rather than string literals for type conversion.
     This function maps common dtype strings to their MLX equivalents.
 
+    float64 survives when the default device supports it and is narrowed to float32
+    otherwise. complex128 is always narrowed, MLX having no wider complex type.
+
     Parameters
     ----------
     dtype_str : str or MLX dtype
         The dtype to convert
-    auto_cast_unsupported : bool
-        If True, automatically cast unsupported dtypes to supported ones with warnings
+    auto_cast_unsupported : bool, optional
+        If True, narrow dtypes the current device cannot handle, warning as it does so.
+        If False, return the requested dtype and let any later failure surface. Default
+        True.
 
     Returns
     -------
     MLX dtype object
     """
-    import warnings
-
     if isinstance(dtype_str, str):
         if dtype_str == "bool":
             return mx.bool_
@@ -58,11 +69,12 @@ def convert_dtype_to_mlx(dtype_str, auto_cast_unsupported=True):
         elif dtype_str == "float32":
             return mx.float32
         elif dtype_str == "float64":
-            if auto_cast_unsupported:
+            if auto_cast_unsupported and not float64_supported():
                 warnings.warn(
                     "MLX does not support float64 on GPU. Automatically casting to float32. "
-                    "This may result in reduced precision. To avoid this warning, "
-                    "explicitly use float32 in your code or set floatX='float32' in PyTensor config.",
+                    "This may result in reduced precision. To keep float64, run on the CPU "
+                    "device with mx.set_default_device(mx.cpu); to avoid this warning, use "
+                    "float32 or set floatX='float32' in PyTensor config.",
                     UserWarning,
                     stacklevel=3,
                 )
@@ -116,6 +128,10 @@ def mlx_typify(data, **kwargs):
 
 @mlx_typify.register(np.ndarray)
 def mlx_typify_tensor(data, dtype=None, **kwargs):
+    # mx.array narrows float64 input to float32 unless the dtype is named explicitly,
+    # and it does so on the CPU too, where float64 is perfectly usable
+    if dtype is None and data.dtype == np.float64 and float64_supported():
+        dtype = mx.float64
     return _nan_safe_constant(data, dtype=dtype)
 
 

{pytensor-3.2.4 → 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,
+)

pytensor-3.2.4/pytensor/link/mlx/dispatch/scalar.py → pytensor-3.3.1/pytensor/link/mlx/dispatch/scalar/basic.py

@@ -5,16 +5,15 @@ import mlx.core as mx
 from pytensor.link.mlx.dispatch.basic import convert_dtype_to_mlx, mlx_funcify
 from pytensor.scalar.basic import (
     Cast,
+    Clip,
     Composite,
     Identity,
     Mod,
     ScalarOp,
     Second,
 )
-from pytensor.scalar.math import Erfc, Erfcx, Log1mexp, Sigmoid, Softplus
 
 
-# MLX name overrides for nfunc_spec names that don't match mlx.core
 MLX_NFUNC_OVERRIDES = {
     "true_divide": "divide",
     "invert": "bitwise_invert",
@@ -96,6 +95,14 @@ def mlx_funcify_Mod(op, **kwargs):
     return mlx_mod
 
 
+@mlx_funcify.register(Clip)
+def mlx_funcify_Clip(op, **kwargs):
+    def clip(x, min, max):
+        return mx.where(x < min, min, mx.where(x > max, max, x))
+
+    return clip
+
+
 @mlx_funcify.register(Identity)
 def mlx_funcify_Identity(op, **kwargs):
     def identity(x):
@@ -115,55 +122,6 @@ def mlx_funcify_Second(op, **kwargs):
     return second
 
 
-@mlx_funcify.register(Sigmoid)
-def mlx_funcify_Sigmoid(op, **kwargs):
-    return mx.sigmoid
-
-
-@mlx_funcify.register(Erfc)
-def mlx_funcify_Erfc(op, **kwargs):
-    def erfc(x):
-        return 1.0 - mx.erf(x)
-
-    return erfc
-
-
-@mlx_funcify.register(Erfcx)
-def mlx_funcify_Erfcx(op, **kwargs):
-    def erfcx(x):
-        return mx.exp(x * x) * (1.0 - mx.erf(x))
-
-    return erfcx
-
-
-@mlx_funcify.register(Softplus)
-def mlx_funcify_Softplus(op, **kwargs):
-    def softplus(x):
-        return mx.where(
-            x < -37.0,
-            mx.exp(x),
-            mx.where(
-                x < 18.0,
-                mx.log1p(mx.exp(x)),
-                mx.where(
-                    x < 33.3,
-                    x + mx.exp(-x),
-                    x,
-                ),
-            ),
-        )
-
-    return softplus
-
-
-@mlx_funcify.register(Log1mexp)
-def mlx_funcify_Log1mexp(op, node, **kwargs):
-    def log1mexp(x):
-        return mx.where(x < mx.log(0.5), mx.log1p(-mx.exp(x)), mx.log(-mx.expm1(x)))
-
-    return log1mexp
-
-
 @mlx_funcify.register(Composite)
 def mlx_funcify_Composite(op, node=None, **kwargs):
     return mlx_funcify(op.fgraph, squeeze_output=True)