diff --git a/doc/library/tensor/basic.rst b/doc/library/tensor/basic.rst index 360b2b8c51..5c9bff756a 100644 --- a/doc/library/tensor/basic.rst +++ b/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 `__ +.. 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. diff --git a/pytensor/link/jax/dispatch/scalar.py b/pytensor/link/jax/dispatch/scalar.py index 2b2ec40df8..a0ebb78560 100644 --- a/pytensor/link/jax/dispatch/scalar.py +++ b/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") diff --git a/pytensor/scalar/math.py b/pytensor/scalar/math.py index 4a8c08e9e2..792298be56 100644 --- a/pytensor/scalar/math.py +++ b/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) diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py index 84d0be7ae2..daeb279251 100644 --- a/pytensor/tensor/math.py +++ b/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", diff --git a/pytensor/xtensor/math.py b/pytensor/xtensor/math.py index a592c8b7d2..c88080f423 100644 --- a/pytensor/xtensor/math.py +++ b/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(): ... diff --git a/tests/link/jax/test_scalar.py b/tests/link/jax/test_scalar.py index 7320d56aa7..b049bad3a6 100644 --- a/tests/link/jax/test_scalar.py +++ b/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", [ diff --git a/tests/link/numba/test_scalar.py b/tests/link/numba/test_scalar.py index fdd505fcb3..7623851d20 100644 --- a/tests/link/numba/test_scalar.py +++ b/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", diff --git a/tests/tensor/test_math_scipy.py b/tests/tensor/test_math_scipy.py index c72889b825..b9c68d85f7 100644 --- a/tests/tensor/test_math_scipy.py +++ b/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=(