pytensor 3.2.3__tar.gz → 3.3.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.2.3
+Version: 3.3.0
 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.65.1,>=0.58
+Requires-Dist: numba<=0.66.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.65.1,>=0.58; extra == "numba"
+Requires-Dist: numba<=0.66.0,>=0.58; extra == "numba"
 Requires-Dist: llvmlite; extra == "numba"
 Provides-Extra: kanren
 Requires-Dist: etuples; extra == "kanren"

{pytensor-3.2.3 → pytensor-3.3.0}/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.65.1",
+    "numba>=0.58,<=0.66.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.65.1", "llvmlite"]
+numba = ["numba>=0.58,<=0.66.0", "llvmlite"]
 kanren = [
     "etuples",
     "logical-unification",
@@ -185,11 +186,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.3 → pytensor-3.3.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-07-24T23:30:42+0200",
+ "date": "2026-08-12T10:16:27+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "42a11f84b1656f2e8c42b3ce05cd602fbcc1c6f8",
- "version": "3.2.3"
+ "full-revisionid": "aa1bc7772eb11ed8def2814374aeb5da03da590c",
+ "version": "3.3.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-3.2.3 → pytensor-3.3.0}/pytensor/compile/rewriting.py

@@ -7,7 +7,7 @@ from pytensor.compile.aliasing import (
     add_supervisor_to_fgraph,
     insert_deepcopy,
 )
-from pytensor.compile.builders import OpFromGraph
+from pytensor.compile.builders import OpFromGraph, SymbolicOp
 from pytensor.compile.io import In, Out
 from pytensor.compile.mode import optdb
 from pytensor.graph.basic import Apply, Variable
@@ -134,12 +134,12 @@ def rewrite_ofg_inner_graph(linker, op, node, inner, *, mode):
     )
 
 
-def _ofg_inner_optimizer(mode):
-    # Recognition rewrites fold a pattern into an inner-graph op (e.g.
-    # ``exp(x) / sum(exp(x))`` -> ``Softmax``, itself an ``OpFromGraph``). Running
-    # them on an ``OpFromGraph`` inner graph -- which may *be* that pattern --
-    # would re-create the op inside itself and recurse without end.
-    return mode.excluding("symbolic_op_recognition").optimizer
+def _ofg_inner_optimizer(mode, op):
+    # Recognition would re-create a `SymbolicOp` inside its own inner graph and never
+    # terminate. Any other `OpFromGraph` must keep it, or wrapping destabilizes the body.
+    if isinstance(op, SymbolicOp):
+        return mode.excluding("symbolic_op_recognition").optimizer
+    return mode.optimizer
 
 
 @rewrite_ofg_inner_graph.register(VMLinker)
@@ -152,7 +152,7 @@ def destructive_rewrite_ofg_inner_graph(linker, op, node, inner, *, mode):
     # still be baked between purely internal buffers.
     input_specs = [In(x, borrow=True, mutable=False) for x in inner.inputs]
     add_supervisor_to_fgraph(fgraph=inner, input_specs=input_specs, accept_inplace=True)
-    _ofg_inner_optimizer(mode).rewrite(inner)
+    _ofg_inner_optimizer(mode, op).rewrite(inner)
     # The op's outputs must not alias its inputs or each other (it declares no
     # view_map, so the outer graph cannot see such aliases); deepcopies break any
     # boundary alias the optimized graph ends up with.
@@ -165,7 +165,7 @@ def destructive_rewrite_ofg_inner_graph(linker, op, node, inner, *, mode):
 @rewrite_ofg_inner_graph.register(MLXLinker)
 def functional_rewrite_ofg_inner_graph(linker, op, node, inner, *, mode):
     """Structurally optimize the inner graph for the functional JIT backends."""
-    _ofg_inner_optimizer(mode).rewrite(inner)
+    _ofg_inner_optimizer(mode, op).rewrite(inner)
 
 
 @graph_rewriter

{pytensor-3.2.3 → pytensor-3.3.0}/pytensor/graph/features.py

@@ -546,10 +546,10 @@ class FullHistory(Feature):
                      └─ ···
         >> local_softmax_stabilize
         Log [id A] 1
-         └─ Softmax{axis=None} [id B] 0
+         └─ Softmax{axis=(0,)} [id B] 0
             └─ x [id C]
         >> local_logsoftmax
-        LogSoftmax{axis=None} [id A] 0
+        LogSoftmax{axis=(0,)} [id A] 0
          └─ x [id B]
 
 
@@ -563,7 +563,7 @@ class FullHistory(Feature):
     .. testoutput::
         >> local_logsoftmax
         Log [id A] 1
-         └─ Softmax{axis=None} [id B] 0
+         └─ Softmax{axis=(0,)} [id B] 0
             └─ x [id C]
         >> local_softmax_stabilize
         Log [id A] 4
@@ -591,7 +591,7 @@ class FullHistory(Feature):
 
     .. testoutput::
         Log [id A] 1
-         └─ Softmax{axis=None} [id B] 0
+         └─ Softmax{axis=(0,)} [id B] 0
             └─ x [id C]
 
 

{pytensor-3.2.3 → pytensor-3.3.0}/pytensor/graph/rewriting/unify.py

@@ -572,7 +572,8 @@ def reify_pattern(pattern, subs: Mapping[PatternVar | Asterisk, Any]):
                 inputs.extend(captured)
             else:
                 inputs.append(reify_pattern(p, subs))
-        return op.make_node(*inputs).default_output()
+        # Call the Op, so those that build state lazily (SymbolicOp) are constructed
+        return op(*inputs)
 
     if isinstance(pattern, OpPattern):
         op_type = pattern.op_type

{pytensor-3.2.3 → pytensor-3.3.0}/pytensor/link/jax/dispatch/scalar.py

@@ -263,14 +263,29 @@ def jax_funcify_Erfinv(op, **kwargs):
 
 
 @jax_funcify.register(BetaIncInv)
-@jax_funcify.register(Erfcx)
-@jax_funcify.register(Erfcinv)
 def jax_funcify_from_tfp(op, **kwargs):
     tfp_jax_op = try_import_tfp_jax_op(op)
 
     return tfp_jax_op
 
 
+@jax_funcify.register(Erfcx)
+def jax_funcify_Erfcx(op, **kwargs):
+    if hasattr(jax.scipy.special, "erfcx"):
+        return jax.scipy.special.erfcx
+    # jax < 0.11 has no native erfcx
+    return try_import_tfp_jax_op(op)
+
+
+@jax_funcify.register(Erfcinv)
+def jax_funcify_Erfcinv(op, **kwargs):
+    def erfcinv(x):
+        # erfc(z) = 2 ndtr(-z * sqrt(2)), so z = -ndtri(x / 2) / sqrt(2)
+        return -jax.scipy.special.ndtri(x / 2) / jnp.sqrt(2)
+
+    return erfcinv
+
+
 @jax_funcify.register(NdtriExp)
 def jax_funcify_NdtriExp(op, **kwargs):
     def ndtri_exp(x):

{pytensor-3.2.3 → pytensor-3.3.0}/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.3 → pytensor-3.3.0}/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.3.0/pytensor/link/mlx/dispatch/blockwise.py

@@ -0,0 +1,64 @@
+import mlx.core as mx
+
+from pytensor.link.mlx.dispatch import mlx_funcify
+from pytensor.tensor.blockwise import Blockwise, _check_runtime_broadcast_core
+
+
+@mlx_funcify.register(Blockwise)
+def funcify_Blockwise(op: Blockwise, node, **kwargs):
+    core_node = op._create_dummy_core_node(node.inputs)
+    core_f = mlx_funcify(op.core_op, node=core_node)
+
+    batch_ndim = op.batch_ndim(node)
+    if batch_ndim == 0:
+        return core_f
+
+    multi_output = len(node.outputs) > 1
+    core_ndims = [len(sig) for sig in op.inputs_sig]
+
+    # Hoisted out of the per-call path, unlike Blockwise._check_runtime_broadcast.
+    batch_bcast = [inp.type.broadcastable[:batch_ndim] for inp in node.inputs]
+
+    # Decide batching purely from static shapes so a graph batches identically
+    # here and in every other backend: a batch axis broadcasts (is never mapped)
+    # only when its static size is exactly 1, or the input lacks it entirely.
+    squeeze_axes, padded_batch = [], []
+    for inp, n_core_dims in zip(node.inputs, core_ndims):
+        batch_shape = inp.type.shape[: inp.type.ndim - n_core_dims]
+        squeeze_axes.append(tuple(i for i, s in enumerate(batch_shape) if s == 1))
+        padded_batch.append((1,) * (batch_ndim - len(batch_shape)) + tuple(batch_shape))
+
+    # Nest one mx.vmap per mapped batch axis (innermost first, so array axis 0
+    # tracks the outermost batch dim). All-broadcast axes are squeezed out of
+    # every input above and re-inserted as size-1 dims after the mapped call.
+    fn, expand_axes = core_f, []
+    for axis in reversed(range(batch_ndim)):
+        in_axes = tuple(None if shape[axis] == 1 else 0 for shape in padded_batch)
+        if all(ax is None for ax in in_axes):
+            expand_axes.append(axis)
+        else:
+            fn = mx.vmap(fn, in_axes=in_axes)
+
+    expand_axes.sort()
+
+    def blockwise(*args):
+        # Other backends reject a runtime size-1 batch dim that is not statically
+        # broadcastable; match them rather than silently broadcasting here.
+        _check_runtime_broadcast_core(args, batch_bcast, batch_ndim)
+
+        squeezed = [
+            mx.squeeze(arg, axes) if axes else arg
+            for arg, axes in zip(args, squeeze_axes)
+        ]
+        out = fn(*squeezed)
+        if not expand_axes:
+            return tuple(out) if multi_output else out
+
+        # Re-insert the never-mapped all-broadcast axes as size-1 dims, in
+        # ascending order so each insertion's index stays valid for the next.
+        outs = out if multi_output else (out,)
+        for ax in expand_axes:
+            outs = [mx.expand_dims(o, ax) for o in outs]
+        return tuple(outs) if multi_output else outs[0]
+
+    return blockwise

{pytensor-3.2.3 → pytensor-3.3.0}/pytensor/link/mlx/dispatch/linalg/decomposition.py

@@ -139,14 +139,15 @@ def mlx_funcify_PivotToPermutations(op, **kwargs):
     inverse = op.inverse
 
     def pivot_to_permutations(pivots):
-        pivots = mx.array(pivots)
-        n = pivots.shape[0]
-        p_inv = mx.arange(n, dtype=mx.int32)
-        for i in range(n):
-            p_inv[i], p_inv[pivots[i]] = p_inv[pivots[i]], p_inv[i]
-        if inverse:
-            return p_inv
-        return mx.argsort(p_inv)
+        with mx.stream(mx.cpu):
+            pivots = mx.array(pivots)
+            n = pivots.shape[0]
+            p_inv = mx.arange(n, dtype=mx.int32)
+            for i in range(n):
+                p_inv[i], p_inv[pivots[i]] = p_inv[pivots[i]], p_inv[i]
+            if inverse:
+                return p_inv
+            return mx.argsort(p_inv)
 
     return pivot_to_permutations
 

pytensor-3.3.0/pytensor/link/mlx/dispatch/linalg/summary.py

@@ -0,0 +1,38 @@
+import mlx.core as mx
+
+from pytensor.link.mlx.dispatch.basic import mlx_funcify
+from pytensor.tensor.linalg.summary import Det, SLogDet
+
+
+def _lu_det_parts(x):
+    """Compute sign and logdet via LU factorization. Call within a CPU stream context."""
+    lu, pivots = mx.linalg.lu_factor(x)
+    diag_u = mx.diagonal(lu)
+    n_swaps = mx.sum(pivots != mx.arange(pivots.shape[0], dtype=pivots.dtype))
+    pivot_sign = 1 - 2 * (n_swaps % 2)
+    sign = pivot_sign * mx.prod(mx.sign(diag_u))
+    logabsdet = mx.sum(mx.log(mx.abs(diag_u)))
+    return sign, logabsdet
+
+
+@mlx_funcify.register(Det)
+def mlx_funcify_Det(op, node, **kwargs):
+    X_dtype = getattr(mx, node.inputs[0].dtype)
+
+    def det(x):
+        with mx.stream(mx.cpu):
+            sign, logabsdet = _lu_det_parts(x.astype(dtype=X_dtype))
+            return sign * mx.exp(logabsdet)
+
+    return det
+
+
+@mlx_funcify.register(SLogDet)
+def mlx_funcify_SLogDet(op, node, **kwargs):
+    X_dtype = getattr(mx, node.inputs[0].dtype)
+
+    def slogdet(x):
+        with mx.stream(mx.cpu):
+            return _lu_det_parts(x.astype(dtype=X_dtype))
+
+    return slogdet

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

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

pytensor-3.2.3/pytensor/link/mlx/dispatch/scalar.py → pytensor-3.3.0/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)