pytensor 3.2.1__tar.gz → 3.2.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.2.1
+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.1 → 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.1 → pytensor-3.2.2}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-07-17T16:59:23+0200",
+ "date": "2026-07-19T19:39:04+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "491faffa05d1e92419dba8e02b245eb2b6227a60",
- "version": "3.2.1"
+ "full-revisionid": "301ed014903acb2d1ed1a34466819cf22f8a3a2b",
+ "version": "3.2.2"
 }
 '''  # END VERSION_JSON
 

{pytensor-3.2.1 → 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.1 → 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.1 → 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.1 → 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.1 → 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.1 → 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.1 → 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.1 → 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.1 → pytensor-3.2.2/pytensor.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.2.1
+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