Add ndtri_exp Op

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/link/jax/dispatch/scalar.py

@@ -32,6 +32,7 @@
     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/scalar/math.py

@@ -268,6 +268,40 @@ def c_code(self, node, name, inp, out, sub):
 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/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 @@ def nan_to_num(x, nan=0.0, posinf=None, neginf=None):
     "mod",
     "mul",
     "nan_to_num",
+    "ndtri_exp",
     "neg",
     "neq",
     "not_equal",

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(): ...
 

tests/link/jax/test_scalar.py

@@ -23,6 +23,7 @@
     kve,
     log,
     log1mexp,
+    ndtri_exp,
     polygamma,
     psi,
     sigmoid,
@@ -143,6 +144,13 @@ def test_erfinv():
     compare_jax_and_py([x], [out], [0.95])
 
 
+def test_ndtri_exp():
+    x = scalar("x")
+    out = ndtri_exp(x)
+
+    compare_jax_and_py([x], [out], [-4.0])
+
+
 @pytest.mark.parametrize(
     "op, test_values",
     [

tests/link/numba/test_scalar.py

@@ -247,6 +247,7 @@ def test_erf_complex():
     (pt.erfcx, [np.array([-1.0, 0.0, 1.0, 3.0])]),
     (pt.erfinv, [np.array([-0.5, 0.0, 0.3, 0.9])]),
     (pt.erfcinv, [np.array([0.2, 0.7, 1.0, 1.5])]),
+    (pt.ndtri_exp, [np.array([-745.0, -100.0, -5.0, -0.5])]),
     (pt.psi, [np.array([0.5, 1.0, 3.7, 12.3])]),
     (pt.gamma, [np.array([0.5, 1.0, 3.7, 5.2])]),
     (pt.j0, [np.array([0.1, 1.0, 5.0, 9.0])]),
@@ -282,6 +283,7 @@ def test_erf_complex():
         "erfcx",
         "erfinv",
         "erfcinv",
+        "ndtri_exp",
         "psi",
         "gamma",
         "j0",

tests/tensor/test_math_scipy.py

@@ -52,6 +52,7 @@ def scipy_special_gammal(k, x):
 expected_erfc = special.erfc
 expected_erfinv = special.erfinv
 expected_erfcinv = special.erfcinv
+expected_ndtri_exp = special.ndtri_exp
 expected_owenst = special.owens_t
 expected_gamma = special.gamma
 expected_gammaln = special.gammaln
@@ -127,6 +128,25 @@ def scipy_special_gammal(k, x):
     mode=mode_no_scipy,
 )
 
+TestNdtriExpBroadcast = makeBroadcastTester(
+    op=pt.ndtri_exp,
+    expected=expected_ndtri_exp,
+    good={
+        "normal": [random_ranged(-10, -0.1, (2, 3))],
+        # exp(x) underflows here, so these are only accurate when computed
+        # directly in logspace
+        "extreme": [
+            np.array(
+                [[-50.0, -100.0, -500.0], [-745.0, -1000.0, -1e6]], dtype=config.floatX
+            )
+        ],
+        "empty": [np.asarray([], dtype=config.floatX)],
+    },
+    grad={"normal": [random_ranged(-5, -0.5, (2, 3))]},
+    eps=2e-10,
+    mode=mode_no_scipy,
+)
+
 rng = np.random.default_rng(seed=utt.fetch_seed())
 _good_broadcast_binary_owenst = dict(
     normal=(