pytensor 3.2.2__tar.gz → 3.2.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.2.2
+Version: 3.2.4
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause
@@ -30,7 +30,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.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 +56,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.2 → pytensor-3.2.4}/pyproject.toml

@@ -49,7 +49,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 +77,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",

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

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-07-19T19:39:04+0200",
+ "date": "2026-08-02T00:12:33+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "301ed014903acb2d1ed1a34466819cf22f8a3a2b",
- "version": "3.2.2"
+ "full-revisionid": "7175f7cadf3934a03f63e6e7b187fa1519e6d85f",
+ "version": "3.2.4"
 }
 '''  # END VERSION_JSON
 

{pytensor-3.2.2 → pytensor-3.2.4}/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.2 → pytensor-3.2.4}/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.2 → pytensor-3.2.4}/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.2 → pytensor-3.2.4}/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.4/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.2 → pytensor-3.2.4}/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.2.4/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.2.2 → pytensor-3.2.4}/pytensor/link/numba/dispatch/elemwise.py

@@ -993,7 +993,7 @@ def numba_funcify_FusedElemwise(op, node, **kwargs):
         def ov_fused_elemwise_fn(*outer_inputs):
             return impl_fn
 
-    cache_version = 3
+    cache_version = 4
     if scalar_cache_key is None:
         key = None
     else:

{pytensor-3.2.2 → pytensor-3.2.4}/pytensor/link/numba/dispatch/random.py

@@ -137,7 +137,7 @@ def numba_core_HalfNormalRV(op, node):
 def numba_core_CauchyRV(op, node):
     @numba_basic.numba_njit
     def random(rng, loc, scale):
-        return (loc + rng.standard_cauchy()) / scale
+        return loc + scale * rng.standard_cauchy()
 
     return random
 

{pytensor-3.2.2 → pytensor-3.2.4}/pytensor/link/numba/dispatch/vectorize_codegen.py

@@ -167,10 +167,20 @@ def _compute_idx_load_axes(indexed_inputs, indexed_outputs, idx_ndims):
             _, out_axis, _ = entry
             group_min_axis[root] = min(group_min_axis.get(root, out_axis), out_axis)
 
-    return tuple(
-        tuple(range(group_min_axis[find(k)], group_min_axis[find(k)] + idx_ndims[k]))
-        for k in range(n_indices)
-    )
+    # Indices in a group broadcast together right-aligned (NumPy semantics), so
+    # a lower-ndim index occupies the trailing loop dims of the group's span,
+    # whose width is the max ndim across the group.
+    group_max_ndim: dict[int, int] = {}
+    for k in range(n_indices):
+        root = find(k)
+        group_max_ndim[root] = max(group_max_ndim.get(root, 0), idx_ndims[k])
+
+    load_axes = []
+    for k in range(n_indices):
+        root = find(k)
+        end = group_min_axis[root] + group_max_ndim[root]
+        load_axes.append(tuple(range(end - idx_ndims[k], end)))
+    return tuple(load_axes)
 
 
 def _core_slice_layout(full_layout, batch_bc_pattern):

{pytensor-3.2.2 → pytensor-3.2.4}/pytensor/scalar/math.py

@@ -298,6 +298,9 @@ class NdtriExp(UnaryScalarOp):
         cst = np.asarray(np.sqrt(2 * np.pi), dtype=gz.type.dtype)
         return (gz * cst * exp(x + z**2 / 2),)
 
+    def c_code(self, *args, **kwargs):
+        raise NotImplementedError()
+
 
 ndtri_exp = NdtriExp(upgrade_to_float_no_complex, name="ndtri_exp")
 

{pytensor-3.2.2 → pytensor-3.2.4}/pytensor/tensor/__init__.py

@@ -126,6 +126,10 @@ from pytensor.tensor.interpolate import interp, interpolate1d
 from pytensor.tensor.math import *
 from pytensor.tensor.pad import pad
 
+# Kept in the top-level namespace for backwards compatibility; the rest of
+# `pytensor.tensor.special` is only reachable as `pt.special.*`
+from pytensor.tensor.special import logaddexp, logsumexp
+
 
 # isort: off
 # reshape needs to be imported before shape.reshape, otherwise the tensor.reshape imports fail

{pytensor-3.2.2 → pytensor-3.2.4}/pytensor/tensor/math.py

@@ -3891,53 +3891,6 @@ def power(x, y):
     return x**y
 
 
-def logaddexp(*xs):
-    """Logarithm of the sum of exponentiations of the inputs.
-
-    See ``numpy.logaddexp``.
-
-    Parameters
-    ----------
-    xs : symbolic tensors
-        Input
-
-    Returns
-    -------
-    TensorVariable
-
-    """
-
-    return log(add(*[exp(x) for x in xs]))
-
-
-def logsumexp(x, axis=None, keepdims=False):
-    """Compute the log of the sum of exponentials of input elements.
-
-    See ``scipy.special.logsumexp``.
-
-    Parameters
-    ----------
-    x : symbolic tensor
-        Input
-
-    axis : None or int or tuple of ints, optional
-        Axis or axes over which the sum is taken. By default axis is None,
-        and all elements are summed.
-
-    keepdims : bool, optional
-        If this is set to True, the axes which are reduced are left in the
-        result as dimensions with size one. With this option, the result will
-        broadcast correctly against the original array.
-
-    Returns
-    -------
-    TensorVariable
-
-    """
-
-    return log(sum(exp(x), axis=axis, keepdims=keepdims))
-
-
 _matmul = Blockwise(_dot, name="Matmul")
 
 
@@ -4321,8 +4274,6 @@ __all__ = [
     "log1pexp",
     "log2",
     "log10",
-    "logaddexp",
-    "logsumexp",
     "lt",
     "matmul",
     "matvec",

{pytensor-3.2.2 → pytensor-3.2.4}/pytensor/tensor/rewriting/math.py

@@ -71,7 +71,6 @@ from pytensor.tensor.math import (
     expm1,
     ge,
     int_div,
-    isinf,
     ive,
     kve,
     le,
@@ -79,8 +78,6 @@ from pytensor.tensor.math import (
     log1mexp,
     log1p,
     log1pexp,
-    makeKeepDims,
-    maximum,
     mul,
     neg,
     polygamma,
@@ -100,9 +97,7 @@ from pytensor.tensor.math import (
     variadic_mul,
 )
 from pytensor.tensor.math import abs as pt_abs
-from pytensor.tensor.math import max as pt_max
 from pytensor.tensor.math import pow as pt_pow
-from pytensor.tensor.math import sum as pt_sum
 from pytensor.tensor.rewriting.basic import (
     broadcast_like_elemwise,
     local_second_sink,
@@ -2813,87 +2808,6 @@ def local_log1p(fgraph, node):
         return [broadcast_like_elemwise(new_out, node, fgraph=fgraph, stack_trace=True)]
 
 
-@register_stabilize("fast_compile")
-@register_specialize
-@node_rewriter([log])
-def local_log_add_exp(fgraph, node):
-    """
-    ``log(exp(x)+exp(y)+exp(z)) = max + log(x-max, y-max, z-max)``
-
-    TODO: in canonicalize, change log10 and log2 -> log
-    """
-
-    z = node.inputs[0]
-    if z.owner and z.owner.op == add:
-        zi = z.owner.inputs
-        pre_exp = [x.owner.inputs[0] for x in zi if x.owner and x.owner.op == exp]
-        # all arguments to add are exp(<something>)
-        if len(pre_exp) == len(zi):
-            # Do not offset when max_pre = -np.inf, to avoid nan in the output
-            # Switch statement is placed directly inside add to break the self-symmetry
-            # of the returned output (otherwise the rewrite would not stabilize)
-            max_pre = reduce(maximum, pre_exp)
-            ret = max_pre + log(
-                add(
-                    *[
-                        switch(isinf(max_pre), exp(max_pre), exp(p - max_pre))
-                        for p in pre_exp
-                    ]
-                )
-            )
-            return [ret]
-
-
-@register_stabilize("fast_compile")
-@register_specialize
-@node_rewriter([log])
-def local_log_sum_exp(fgraph, node):
-    # log(sum_i(exp(x_i))) = x_max + log(sum_i(exp(x_i - x_max)))
-
-    sum_node = node.inputs[0].owner
-    # If the sum has keepdims=True, there might be a dimshuffle
-    if sum_node and isinstance(sum_node.op, DimShuffle):
-        dimshuffle_op = sum_node.op
-        sum_node = sum_node.inputs[0].owner
-    else:
-        dimshuffle_op = None
-
-    if not (sum_node and isinstance(sum_node.op, Sum)):
-        return
-
-    exp_node, axis = sum_node.inputs[0].owner, sum_node.op.axis
-    if not (
-        exp_node
-        and isinstance(exp_node.op, Elemwise)
-        and isinstance(exp_node.op.scalar_op, ps.Exp)
-    ):
-        return
-
-    pre_exp = exp_node.inputs[0]
-    max_pre_exp = pt_max(pre_exp, axis=axis)
-    max_pre_exp_keepdims = makeKeepDims(pre_exp, max_pre_exp, axis)
-
-    # Do not offset when max_pre = -np.inf, to avoid nan in the output
-    # Switch statement is placed directly inside sum to break the self-symmetry
-    # of the returned output (otherwise the rewrite would not stabilize)
-    ret = max_pre_exp + log(
-        pt_sum(
-            switch(
-                isinf(max_pre_exp_keepdims),
-                exp(max_pre_exp_keepdims),
-                exp(pre_exp - max_pre_exp_keepdims),
-            ),
-            axis=axis,
-        ),
-    )
-
-    # Restore the dimshuffle op, if any.
-    if dimshuffle_op:
-        ret = dimshuffle_op(ret)
-
-    return [ret]
-
-
 def add_calculate(num, denum, aslist=False, out_type=None):
     # TODO: make sure that this function and mul_calculate are similar
     if out_type is None:

{pytensor-3.2.2 → pytensor-3.2.4}/pytensor/tensor/rewriting/ofg.py

@@ -2,11 +2,11 @@ from pytensor.compile.rewriting import inline_ofg_node
 from pytensor.graph import node_rewriter
 from pytensor.tensor.basic import AllocDiag
 from pytensor.tensor.rewriting.basic import register_specialize
-from pytensor.tensor.special import XLog1PY, XLogY
+from pytensor.tensor.special import LogAddExp, LogSumExp, XLog1PY, XLogY
 
 
 @register_specialize("inline_ofg")
-@node_rewriter([AllocDiag, XLogY, XLog1PY])
+@node_rewriter([AllocDiag, XLogY, XLog1PY, LogSumExp, LogAddExp])
 def late_inline_OpFromGraph(fgraph, node):
     """
     Inline `OpFromGraph` nodes.