pytensor 3.2.0__tar.gz → 3.2.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.2.0
+Version: 3.2.2
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause

{pytensor-3.2.0 → pytensor-3.2.2}/doc/library/tensor/basic.rst

@@ -1528,6 +1528,10 @@ Mathematical
 
     Returns a variable representing the inverse error function or the inverse complementary error function. `wikipedia <http://en.wikipedia.org/wiki/Error_function#Inverse_functions>`__
 
+.. function:: ndtri_exp(a)
+
+    Returns a variable representing the inverse of the standard normal CDF evaluated at the exponent of a, computed accurately even where ``exp(a)`` underflows.
+
 .. function:: gamma(a)
 
    Returns a variable representing the gamma function.

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

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-07-16T16:05:00+0200",
+ "date": "2026-07-19T19:39:04+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "5480dc534c09f11c8cc5b0cf66374eabb44f0c6e",
- "version": "3.2.0"
+ "full-revisionid": "301ed014903acb2d1ed1a34466819cf22f8a3a2b",
+ "version": "3.2.2"
 }
 '''  # END VERSION_JSON
 

{pytensor-3.2.0 → pytensor-3.2.2}/pytensor/link/jax/dispatch/scalar.py

@@ -32,6 +32,7 @@ from pytensor.scalar.math import (
     Ive,
     Kve,
     Log1mexp,
+    NdtriExp,
     Psi,
     TriGamma,
 )
@@ -270,6 +271,16 @@ def jax_funcify_from_tfp(op, **kwargs):
     return tfp_jax_op
 
 
+@jax_funcify.register(NdtriExp)
+def jax_funcify_NdtriExp(op, **kwargs):
+    def ndtri_exp(x):
+        # JAX has no ndtri_exp, so this composition loses accuracy where
+        # exp(x) underflows
+        return jax.scipy.special.ndtri(jnp.exp(x))
+
+    return ndtri_exp
+
+
 @jax_funcify.register(Ive)
 def jax_funcify_Ive(op, **kwargs):
     return try_import_tfp_jax_op(op, jax_op_name="bessel_ive")

{pytensor-3.2.0 → pytensor-3.2.2}/pytensor/link/jax/dispatch/subtensor.py

@@ -59,8 +59,8 @@ def jax_funcify_IncSubtensor(op, node, **kwargs):
         if len(indices) == 1:
             indices = indices[0]
 
-        if isinstance(op, AdvancedIncSubtensor) and op.idx_list == (0,):
-            op._check_runtime_broadcasting(node, x, y, indices)
+        if isinstance(op, AdvancedIncSubtensor):
+            op._check_runtime_broadcast_of_vector_index(node, x, y, indices)
 
         return jax_fn(x, indices, y)
 

{pytensor-3.2.0 → pytensor-3.2.2}/pytensor/link/jax/dispatch/tensor_basic.py

@@ -71,7 +71,10 @@ def jax_funcify_ARange(op, node, **kwargs):
             case (Subtensor(), shape_var, *_) if isinstance(shape_var.owner_op, Shape):
                 constant_args.append(None)
             case _ if isinstance(arg, Constant):
-                constant_args.append(arg.value)
+                # Cast to the Op's dtype: PyTensor types integer literals (e.g. 0/1 arange
+                # start/step) as int8, and jnp.arange bounds-checks stop against the argument dtype,
+                # overflowing for stop > 127.
+                constant_args.append(np.asarray(arg.value, op.dtype))
             case _:
                 # TODO: This might be failing without need (e.g., if arg = shape(x)[-1] + 1)!
                 raise NotImplementedError(ARANGE_CONCRETE_VALUE_ERROR)

{pytensor-3.2.0 → pytensor-3.2.2}/pytensor/link/mlx/dispatch/subtensor.py

@@ -97,8 +97,7 @@ def mlx_funcify_AdvancedIncSubtensor(op, node, **kwargs):
             return x.at[indices].add(y)
 
     def advancedincsubtensor(x, y, *ilist, mlx_fn=mlx_fn):
-        if op.idx_list == (0,):
-            op._check_runtime_broadcasting(node, x, y, ilist[0])
+        op._check_runtime_broadcast_of_vector_index(node, x, y, ilist[0])
 
         return mlx_fn(x, ilist, y)
 

{pytensor-3.2.0 → pytensor-3.2.2}/pytensor/link/pytorch/dispatch/subtensor.py

@@ -63,8 +63,6 @@ def pytorch_funcify_IncSubtensor(op, node, **kwargs):
         def set_subtensor(x, y, *flattened_indices):
             indices = indices_from_subtensor(flattened_indices, idx_list)
             check_negative_steps(indices)
-            if op.idx_list == (0,):
-                op._check_runtime_broadcasting(node, x, y, indices[0])
             if not inplace:
                 x = x.clone()
             x[indices] = y
@@ -77,8 +75,6 @@ def pytorch_funcify_IncSubtensor(op, node, **kwargs):
         def inc_subtensor(x, y, *flattened_indices):
             indices = indices_from_subtensor(flattened_indices, idx_list)
             check_negative_steps(indices)
-            if op.idx_list == (0,):
-                op._check_runtime_broadcasting(node, x, y, indices[0])
             if not inplace:
                 x = x.clone()
             x[indices] += y
@@ -98,8 +94,7 @@ def pytorch_funcify_AdvancedIncSubtensor(op, node, **kwargs):
         def adv_set_subtensor(x, y, *flattened_indices):
             indices = indices_from_subtensor(flattened_indices, idx_list)
             check_negative_steps(indices)
-            if op.idx_list == (0,):
-                op._check_runtime_broadcasting(node, x, y, indices[0])
+            op._check_runtime_broadcast_of_vector_index(node, x, y, indices[0])
             if not inplace:
                 x = x.clone()
             x[indices] = y.type_as(x)
@@ -112,8 +107,7 @@ def pytorch_funcify_AdvancedIncSubtensor(op, node, **kwargs):
         def adv_inc_subtensor_no_duplicates(x, y, *flattened_indices):
             indices = indices_from_subtensor(flattened_indices, idx_list)
             check_negative_steps(indices)
-            if op.idx_list == (0,):
-                op._check_runtime_broadcasting(node, x, y, indices[0])
+            op._check_runtime_broadcast_of_vector_index(node, x, y, indices[0])
             if not inplace:
                 x = x.clone()
             x[indices] += y.type_as(x)
@@ -131,8 +125,7 @@ def pytorch_funcify_AdvancedIncSubtensor(op, node, **kwargs):
             indices = indices_from_subtensor(flattened_indices, idx_list)
             # Not needed because slices aren't supported in this path
             # check_negative_steps(indices)
-            if op.idx_list == (0,):
-                op._check_runtime_broadcasting(node, x, y, indices[0])
+            op._check_runtime_broadcast_of_vector_index(node, x, y, indices[0])
             if not inplace:
                 x = x.clone()
             x.index_put_(indices, y.type_as(x), accumulate=True)

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

@@ -268,6 +268,40 @@ class Erfcinv(UnaryScalarOp):
 erfcinv = Erfcinv(upgrade_to_float_no_complex, name="erfcinv")
 
 
+class NdtriExp(UnaryScalarOp):
+    """
+    Implements the inverse of the standard normal CDF evaluated at the
+    exponent of x, `ndtri(exp(x))`, in a way that remains accurate for very
+    negative x, where `exp(x)` underflows.
+    """
+
+    monotonic_increasing = True
+    nfunc_spec = ("scipy.special.ndtri_exp", 1, 1)
+
+    def impl(self, x):
+        return special.ndtri_exp(x)
+
+    def pullback(self, inputs, outputs, grads):
+        (x,) = inputs
+        (z,) = outputs
+        (gz,) = grads
+        if x.type in complex_types:
+            raise NotImplementedError()
+        if z.type in discrete_types:
+            if x.type in discrete_types:
+                return [x.zeros_like(dtype=config.floatX)]
+            else:
+                return [x.zeros_like()]
+
+        # d/dx ndtri(exp(x)) = exp(x) / pdf(z), evaluated as sqrt(2 * pi) * exp(x + z ** 2 / 2)
+        # so that the underflowing exp(x) and the overflowing 1 / pdf(z) never appear on their own
+        cst = np.asarray(np.sqrt(2 * np.pi), dtype=gz.type.dtype)
+        return (gz * cst * exp(x + z**2 / 2),)
+
+
+ndtri_exp = NdtriExp(upgrade_to_float_no_complex, name="ndtri_exp")
+
+
 class Owens_t(BinaryScalarOp):
     nfunc_spec = ("scipy.special.owens_t", 2, 1)
 

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

@@ -2372,6 +2372,11 @@ def erfcinv(a):
     """inverse complementary error function"""
 
 
+@scalar_elemwise
+def ndtri_exp(a):
+    """inverse standard normal cdf of the exponent of a"""
+
+
 @scalar_elemwise
 def owens_t(h, a):
     """owens t function"""
@@ -4331,6 +4336,7 @@ __all__ = [
     "mod",
     "mul",
     "nan_to_num",
+    "ndtri_exp",
     "neg",
     "neq",
     "not_equal",

{pytensor-3.2.0 → pytensor-3.2.2}/pytensor/tensor/rewriting/subtensor.py

@@ -76,7 +76,6 @@ from pytensor.tensor.subtensor import (
     _is_provably_non_negative,
     _is_provably_positive,
     _non_consecutive_adv_indexing,
-    advanced_inc_subtensor,
     as_index_literal,
     flatten_index_variables,
     get_canonical_form_slice,
@@ -1104,11 +1103,13 @@ def local_useless_inc_subtensor(fgraph, node):
 @register_specialize
 @node_rewriter([AdvancedIncSubtensor])
 def local_set_to_inc_subtensor(fgraph, node):
-    r"""
-    AdvancedIncSubtensor(x, x[ilist]+other, ilist, set_instead_of_inc=True) ->
-    AdvancedIncSubtensor(x, other, ilist, set_instead_of_inc=False)
+    r"""Set of a read at the same indices: ``x[idx].set(x[idx] + other)``.
 
-    Only valid when ``ilist`` is duplicate-free: a dense set is last-wins while an
+    .. code::
+
+        x[idx].set(x[idx] + other) -> x[idx].inc(other)
+
+    Only valid when ``idx`` is duplicate-free: a dense set is last-wins while an
     inc accumulates, so duplicate indices would over-count.
 
     TODO FIXME: Why doesn't this apply to all `*IncSubtensor*` `Op`\s?  If it
@@ -1136,14 +1137,30 @@ def local_set_to_inc_subtensor(fgraph, node):
         other = addn.inputs[0]
     else:
         return
-    if subn.inputs[1] != node.inputs[2] or subn.inputs[0] != node.inputs[0]:
+    # The read has to gather exactly the positions the write stores to: same base,
+    # same axes, same indices. `idx_list` must be compared too -- it is what pins
+    # the indices to their axes, so without it a read of ``x[:, i]`` would match a
+    # write to ``x[i]``.
+    if (
+        subn.inputs[0] != node.inputs[0]
+        or subn.op.idx_list != node.op.idx_list
+        or subn.inputs[1:] != node.inputs[2:]
+    ):
         return
-    # set->inc is only valid when ilist is duplicate-free: ``set(x[ilist] + other)``
-    # is last-wins at repeated positions, while ``inc(other)`` would accumulate the
-    # contributions of every occurrence and over-count them.
-    if not _index_provably_unique(node.inputs[2], fgraph):
+    # set->inc is only valid when the written positions are duplicate-free:
+    # ``set(x[idx] + other)`` is last-wins at repeated positions, while
+    # ``inc(other)`` would accumulate the contributions of every occurrence and
+    # over-count them.
+    indices = indices_from_subtensor(node.inputs[2:], node.op.idx_list)
+    if not _advanced_indices_jointly_unique(indices, fgraph):
         return
-    ret = advanced_inc_subtensor(node.inputs[0], other, node.inputs[2])
+    new_op = type(node.op)(
+        idx_list=node.op.idx_list,
+        inplace=node.op.inplace,
+        set_instead_of_inc=False,
+        ignore_duplicates=node.op.ignore_duplicates,
+    )
+    ret = new_op(node.inputs[0], other, *node.inputs[2:])
 
     copy_stack_trace(node.outputs, ret)
 
@@ -1341,52 +1358,6 @@ def local_convert_negative_indices(fgraph, node):
     return [new_subtensor]
 
 
-@register_canonicalize
-@register_specialize
-@node_rewriter([AdvancedSubtensor])
-def local_useless_AdvancedSubtensor1(fgraph, node):
-    """Remove a leading-axis vector take that selects the full input.
-
-    The full input is taken when the index is equivalent to
-    ``arange(0, input.shape[0], 1)`` using either an explicit list/vector or
-    the `ARange` `Op`.
-
-    """
-    if node.op.idx_list != (0,):
-        # Only the leading integer-vector take (x[ilist])
-        return
-
-    # This optimization needs ShapeOpt and fgraph.shape_feature
-    if not hasattr(fgraph, "shape_feature"):
-        return
-
-    shape_feature = fgraph.shape_feature
-
-    # get length of the indexed tensor along the first axis
-    try:
-        length = get_scalar_constant_value(
-            shape_feature.get_shape(node.inputs[0], 0), only_process_constants=True
-        )
-    except NotScalarConstantError:
-        return False
-
-    # get index (which must be a vector by definition)
-    idx = node.inputs[1]
-
-    # `idx` must be equivalent to [0,1,...,shape[0] - 1] to qualify for
-    # this optimization
-    if isinstance(idx, Constant):
-        idx = idx.value
-        if len(idx) != length:
-            return False
-        if np.any(idx != np.arange(length)):
-            return False
-    else:
-        return False
-
-    return [node.inputs[0]]
-
-
 def _arange_index_to_slice(idx):
     """Return the Python ``slice`` equivalent to ``idx`` if it is a constant or
     symbolic arange with non-negative offset, else ``None``.
@@ -2343,8 +2314,7 @@ def local_write_of_write_same_indices(fgraph, node):
     base, a, *inner_idx_vars = inner_x.owner.inputs
 
     # Same indices: idx_list (slice specs) must match and all index
-    # variables must be identical.  AdvancedIncSubtensor1 has a fixed
-    # class-level idx_list = (0,) so the comparison is trivially true.
+    # variables must be identical.
     if node.op.idx_list != inner_x.owner.op.idx_list:
         return
     if not all(o is i for o, i in zip(outer_idx_vars, inner_idx_vars, strict=True)):

{pytensor-3.2.0 → pytensor-3.2.2}/pytensor/tensor/rewriting/subtensor_lift.py

@@ -1062,6 +1062,11 @@ def local_subtensor_make_vector(fgraph, node):
             except NotScalarConstantError:
                 pass
         elif idx.ndim == 1 and isinstance(idx, Constant):
+            if idx.type.dtype == "bool":
+                # A boolean mask selects the entries it is True at, which is not
+                # what mapping it to positions below would read it as. Leave it
+                # for `bool_idx_to_nonzero` to turn into integer positions.
+                return False
             values = list(map(int, list(idx.value)))
             if len(values) == 1:
                 ret = expand_dims(x.owner.inputs[values[0]], axis=0)

{pytensor-3.2.0 → pytensor-3.2.2}/pytensor/tensor/subtensor.py

@@ -2445,29 +2445,32 @@ class AdvancedIncSubtensor(BaseSubtensor, COp):
     def c_code_cache_version(self):
         return (2,)
 
-    def _check_runtime_broadcasting(self, node, x, y, idx) -> None:
-        """Raise if ``y`` would silently broadcast against the leading-axis update.
+    def _check_runtime_broadcast_of_vector_index(self, node, x, y, idx) -> None:
+        """Raise if ``x[idx] (+)= y`` would broadcast ``y`` at runtime.
 
-        Applies to the leading integer-vector form; ``y`` must already match the
-        indexed shape rather than being broadcast at runtime.
+        Only the single integer vector index form is checked, where the update
+        must have shape ``(idx.shape[0], *x.shape[1:])``. Other index forms
+        (boolean masks, multidimensional or multiple indices) are left to the
+        numpy-like broadcasting of the underlying implementation.
         """
-        if y.ndim > 0:
-            y_pt_bcast = node.inputs[1].broadcastable
-
-            if not y_pt_bcast[0] and y.shape[0] == 1 and y.shape[0] != idx.shape[0]:
-                # Attempting to broadcast with index
-                raise ValueError(self._runtime_broadcast_error_msg)
-            if any(
-                not y_bcast and y_dim == 1 and y_dim != x_dim
-                for y_bcast, y_dim, x_dim in zip(
-                    reversed(y_pt_bcast),
-                    reversed(y.shape),
-                    reversed(x.shape),
-                    strict=False,
-                )
-            ):
-                # Attempting to broadcast with buffer
-                raise ValueError(self._runtime_broadcast_error_msg)
+        if self.idx_list != (0,):
+            return
+        idx_type = node.inputs[2].type
+        if idx_type.ndim != 1 or idx_type.dtype == "bool":
+            return
+
+        y_static_bcast = node.inputs[1].type.broadcastable
+        expected_shape = (idx.shape[0], *x.shape[1:])
+        if any(
+            not y_bcast and y_dim == 1 and expected_dim != 1
+            for y_bcast, y_dim, expected_dim in zip(
+                reversed(y_static_bcast),
+                reversed(y.shape),
+                reversed(expected_shape),
+                strict=False,
+            )
+        ):
+            raise ValueError(self._runtime_broadcast_error_msg)
 
     def __init__(
         self,
@@ -2510,8 +2513,7 @@ class AdvancedIncSubtensor(BaseSubtensor, COp):
     def perform(self, node, inputs, out_):
         x, y, *index_variables = inputs
 
-        if self.idx_list == (0,):
-            self._check_runtime_broadcasting(node, x, y, index_variables[0])
+        self._check_runtime_broadcast_of_vector_index(node, x, y, index_variables[0])
 
         full_indices = unflatten_index_variables(index_variables, self.idx_list)
 

{pytensor-3.2.0 → pytensor-3.2.2}/pytensor/xtensor/math.py

@@ -151,6 +151,10 @@ def erfcx(): ...
 def erfinv(): ...
 
 
+@_as_xelemwise(ps.ndtri_exp)
+def ndtri_exp(): ...
+
+
 @_as_xelemwise(ps.exp)
 def exp(): ...
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.2.0
+Version: 3.2.2
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause