diff --git a/pytensor/compile/builders.py b/pytensor/compile/builders.py index 1bf6d54896..59e79f5aeb 100644 --- a/pytensor/compile/builders.py +++ b/pytensor/compile/builders.py @@ -902,8 +902,36 @@ def __init_subclass__(cls, **kwargs): if "__props__" in cls.__dict__: # MetaType installs props-only __hash__ and __eq__ which ignores the inner graph # override with fgraph-aware version - cls.__hash__ = OpFromGraph.__hash__ - cls.__eq__ = OpFromGraph.__eq__ + cls.__hash__ = SymbolicOp.__hash__ + cls.__eq__ = SymbolicOp.__eq__ + + def _deferred_key(self): + """Identity of an op whose inner graph hasn't been built yet. + + Construction is deferred until the first ``__call__`` (see ``__init__``), so + until then there is no ``fgraph`` to identify the op by. Two deferred ops of the + same type and props will build the same inner graph for any given input types, + so the type and props are what identify them. + """ + props = getattr(type(self), "__props__", ()) + return (type(self), tuple(getattr(self, p) for p in props)) + + def __hash__(self): + if not hasattr(self, "fgraph"): + return hash(self._deferred_key()) + return OpFromGraph.__hash__(self) + + def __eq__(self, other): + if self is other: + return True + if type(self) is not type(other): + return False + self_built = hasattr(self, "fgraph") + if self_built != hasattr(other, "fgraph"): + return False + if not self_built: + return self._deferred_key() == other._deferred_key() + return OpFromGraph.__eq__(self, other) @staticmethod def filter_inputs(*inputs): diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index 2921d8e0c3..07e6a2c4f2 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -737,8 +737,6 @@ def vectorize_scalar_from_tensor(op, node, batch_x): return batch_x.copy().owner -# to be removed as we get the epydoc routine-documenting thing going -# -JB 20080924 def _conversion(real_value: Op, name: str) -> Op: __oplist_tag(real_value, "casting") real_value.__module__ = "tensor.basic" diff --git a/pytensor/tensor/elemwise.py b/pytensor/tensor/elemwise.py index 4f3cc3f738..da985e2518 100644 --- a/pytensor/tensor/elemwise.py +++ b/pytensor/tensor/elemwise.py @@ -461,7 +461,6 @@ def __getstate__(self): d = copy(self.__dict__) d.pop("ufunc") d.pop("nfunc") - d.pop("__epydoc_asRoutine", None) return d def __setstate__(self, d): @@ -1741,9 +1740,6 @@ def construct(symbol): if getattr(symbol, "__doc__"): rval.__doc__ = symbol.__doc__ - # for the meaning of this see the ./epydoc script - # it makes epydoc display rval as if it were a function, not an object - rval.__epydoc_asRoutine = symbol rval.__module__ = symbol.__module__ return rval diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py index bc38f4391e..53f9cc2a09 100644 --- a/pytensor/tensor/math.py +++ b/pytensor/tensor/math.py @@ -38,6 +38,7 @@ scalar_elemwise, ) from pytensor.tensor.shape import shape, specify_shape +from pytensor.tensor.symbolic import TensorSymbolicOp from pytensor.tensor.type import ( DenseTensorType, complex_dtypes, @@ -2358,6 +2359,33 @@ def erfcx(a): """scaled complementary error function""" +class LogErfc(TensorSymbolicOp): + """Compute ``log(erfc(x))`` without underflowing in the tail. + + The naive form returns ``-inf`` once ``erfc(x)`` underflows, around x=26.6 in + float64. ``erfcx(x) = exp(x**2) * erfc(x)`` stays finite there, so the identity + ``log(erfc(x)) == log(erfcx(x)) - x**2`` keeps the tail in log space. + + The two branches are each unusable in the other's half: ``erfcx(x) -> 1`` as + ``x -> 0``, so ``log(erfcx(x))`` loses the value to cancellation (5e-5 relative at + x=1e-12), while ``erf(x)`` saturates to 1 around x=6, sending ``log1p(-erf(x))`` to + ``-inf``. Splitting at 1 leaves both accurate to ~1e-16 at the seam. + """ + + def build_inner_graph(self, x): + return [switch(x < 1.0, log1p(-erf(x)), log(erfcx(x)) - sqr(x))] + + def pullback(self, inputs, outputs, output_grads): + (x,) = inputs + (gz,) = output_grads + # d/dx log(erfc(x)) is -2/sqrt(pi) * exp(-x**2) / erfc(x), a 0/0 quotient once + # both underflow. Substituting erfc(x) = exp(-x**2) * erfcx(x) cancels it to + # -2/(sqrt(pi) * erfcx(x)), which is finite for every x: erfcx overflows to inf + # for x << 0, where the derivative does vanish. + cst = np.array(-2.0 / np.sqrt(np.pi), dtype=x.type.dtype) + return [gz * cst / erfcx(x)] + + @scalar_elemwise def erfinv(a): """inverse error function""" @@ -2497,6 +2525,58 @@ def kn(n, x): return kv(n, x) +# `scipy.special.jn` is an alias of `jv`, not a separate function +jn = jv + + +def i0e(x): + """exponentially scaled modified Bessel function of order 0 + + Matches :func:`scipy.special.i0e`. + """ + return ive(0, x) + + +def i1e(x): + """exponentially scaled modified Bessel function of order 1 + + Matches :func:`scipy.special.i1e`. + """ + return ive(1, x) + + +def k0(x): + """modified Bessel function of the second kind of order 0 + + Matches :func:`scipy.special.k0`. + """ + return kv(0, x) + + +def k1(x): + """modified Bessel function of the second kind of order 1 + + Matches :func:`scipy.special.k1`. + """ + return kv(1, x) + + +def k0e(x): + """exponentially scaled modified Bessel function of the second kind of order 0 + + Matches :func:`scipy.special.k0e`. + """ + return kve(0, x) + + +def k1e(x): + """exponentially scaled modified Bessel function of the second kind of order 1 + + Matches :func:`scipy.special.k1e`. + """ + return kve(1, x) + + @scalar_elemwise def sigmoid(x): """Logistic sigmoid function (1 / (1 + exp(-x)), also known as expit or inverse logit""" @@ -4277,7 +4357,9 @@ def nan_to_num(x, nan=0.0, posinf=None, neginf=None): "gt", "hyp2f1", "i0", + "i0e", "i1", + "i1e", "imag", "int_div", "invert", @@ -4292,7 +4374,12 @@ def nan_to_num(x, nan=0.0, posinf=None, neginf=None): "ive", "j0", "j1", + "jn", "jv", + "k0", + "k0e", + "k1", + "k1e", "kn", "kv", "kve", diff --git a/pytensor/tensor/rewriting/math.py b/pytensor/tensor/rewriting/math.py index e2693394b2..2a6b4763af 100644 --- a/pytensor/tensor/rewriting/math.py +++ b/pytensor/tensor/rewriting/math.py @@ -49,6 +49,7 @@ All, Any, Dot, + LogErfc, Max, Min, Prod, @@ -3163,50 +3164,28 @@ def local_greedy_distributor(fgraph, node): register_specialize(local_erf_neg_minus_one) -@register_stabilize -@register_specialize -@node_rewriter([log]) -def local_log_erfc(fgraph, node): - """Stability rewrite for ``log(erfc(x))``. - - Notes - ----- - log(erfc(x)) => when x>threshold, - -x**2-log(x)-.5*log(pi)+log(1-1/(2*x**2)+3/(4*x**4)-15/(8*x**6)) - for float64: threshold=26.641747557 was chosen with: - [(i,numpy.log(scipy.special.erfc(numpy.asarray([i],dtype='float64')))) - for i in numpy.arange(26.641747557,26.6417475571,.00000000001)] - for float32: threshold=10.0541949, [(i,numpy.log(scipy.special.erfc( - numpy.asarray([i],dtype='float32')))) for i in numpy.arange( - 10.0541948,10.0541951,.0000001)] - """ - - if not (node.inputs[0].owner and node.inputs[0].owner.op == erfc): - return False - - if hasattr(node.tag, "local_log_erfc_applied"): - # We use that flag to don't apply the rewrite recursively - # TODO FIXME: We shouldn't need to use tags for this. - return False - - node.tag.local_log_erfc_applied = True - - x = node.inputs[0].owner.inputs[0] - stab_value = ( - -(x**2) - - log(x) - - 0.5 * log(np.pi) - + log(1 - 1 / (2 * x**2) + 3 / (4 * x**4) - 15 / (8 * x**6)) - ) - - if node.outputs[0].dtype == "float32" or node.outputs[0].dtype == "float16": - threshold = 10.0541949 - elif node.outputs[0].dtype == "float64": - threshold = 26.641747557 - - ret = switch(x < threshold, node.outputs[0], stab_value) - ret.tag.values_eq_approx = values_eq_approx_remove_inf - return [ret] +# log(erfc(x)) => LogErfc(x), which stays finite where erfc(x) underflows to 0. +local_log_erfc = PatternNodeRewriter( + (log, (erfc, "x")), + (LogErfc(), "x"), + name="local_log_erfc", + values_eq_approx=values_eq_approx_remove_inf, +) +register_stabilize(local_log_erfc, "symbolic_op_recognition") +register_specialize(local_log_erfc, "symbolic_op_recognition") + +# log(y * erfc(x)) => log(y) + LogErfc(x). erfc is strictly positive, so only y > 0 is +# meaningful, and there the identity holds for any y, constant or not. A y <= 0 +# propagates nan/-inf through log(y) instead. This is the form a log of a scaled erfc +# arrives in, e.g. ndtr(x) == 0.5 * erfc(-x / sqrt(2)). +local_log_mul_erfc = PatternNodeRewriter( + (log, (mul, "y", (erfc, "x"))), + (add, (log, "y"), (LogErfc(), "x")), + name="local_log_mul_erfc", + values_eq_approx=values_eq_approx_remove_inf, +) +register_stabilize(local_log_mul_erfc, "symbolic_op_recognition") +register_specialize(local_log_mul_erfc, "symbolic_op_recognition") @register_stabilize diff --git a/pytensor/tensor/rewriting/ofg.py b/pytensor/tensor/rewriting/ofg.py index 5f8b865fee..8245096d91 100644 --- a/pytensor/tensor/rewriting/ofg.py +++ b/pytensor/tensor/rewriting/ofg.py @@ -1,12 +1,13 @@ from pytensor.compile.rewriting import inline_ofg_node from pytensor.graph import node_rewriter from pytensor.tensor.basic import AllocDiag +from pytensor.tensor.math import LogErfc from pytensor.tensor.rewriting.basic import register_specialize -from pytensor.tensor.special import XLog1PY, XLogY +from pytensor.tensor.special import LogNdtr, Ndtr, XLog1PY, XLogY @register_specialize("inline_ofg") -@node_rewriter([AllocDiag, XLogY, XLog1PY]) +@node_rewriter([AllocDiag, LogErfc, LogNdtr, Ndtr, XLogY, XLog1PY]) def late_inline_OpFromGraph(fgraph, node): """ Inline `OpFromGraph` nodes. diff --git a/pytensor/tensor/rewriting/special.py b/pytensor/tensor/rewriting/special.py index ab5e9b128b..a6ba5b4613 100644 --- a/pytensor/tensor/rewriting/special.py +++ b/pytensor/tensor/rewriting/special.py @@ -1,9 +1,14 @@ -from pytensor.graph.rewriting.basic import copy_stack_trace, node_rewriter +from pytensor.graph.rewriting.basic import ( + PatternNodeRewriter, + copy_stack_trace, + node_rewriter, +) +from pytensor.graph.rewriting.unify import OpPattern from pytensor.scalar.basic import Exp from pytensor.tensor.elemwise import DimShuffle, Elemwise from pytensor.tensor.math import Sum, log, true_div from pytensor.tensor.rewriting.basic import register_stabilize -from pytensor.tensor.special import Softmax, log_softmax +from pytensor.tensor.special import LogNdtr, Ndtr, Softmax, log_softmax from pytensor.tensor.subtensor import ( AdvancedSubtensor, AdvancedSubtensor1, @@ -19,6 +24,17 @@ ) +# Detect Log(Ndtr(x)) and replace it with LogNdtr(x), which stays finite in both tails. +# Note: only forward pass is affected, as the gradient graph is built before this runs. +local_log_ndtr = PatternNodeRewriter( + (log, (OpPattern(Ndtr), "x")), + (LogNdtr(), "x"), + name="local_log_ndtr", + values_eq_approx=values_eq_approx_remove_inf, +) +register_stabilize(local_log_ndtr) + + @register_stabilize @node_rewriter([log]) def local_logsoftmax(fgraph, node): diff --git a/pytensor/tensor/special.py b/pytensor/tensor/special.py index 2cc1acbf52..66bc691b55 100644 --- a/pytensor/tensor/special.py +++ b/pytensor/tensor/special.py @@ -1,17 +1,30 @@ +import numpy as np from numpy.lib.array_utils import normalize_axis_tuple +from pytensor.configdefaults import config from pytensor.gradient import DisconnectedType, disconnected_type from pytensor.graph.replace import _vectorize_node from pytensor.tensor import as_tensor_variable from pytensor.tensor.elemwise import get_normalized_batch_axes from pytensor.tensor.math import ( eq, + erfc, + erfcinv, + erfcx, exp, + floor, gamma, + gammainc, + gammaincc, + gammainccinv, gammaln, + le, log, log1p, + lt, mul, + softplus, + sqr, sum, switch, ) @@ -133,6 +146,223 @@ def logit(x): return log(x / (1 - x)) +def log_expit(x): + """Logarithm of the logistic sigmoid. + + Matches :func:`scipy.special.log_expit`. Unlike ``log(expit(x))``, this does not + underflow to ``-inf`` for large negative x. + + """ + return -softplus(-x) + + +def rgamma(x): + """Reciprocal of the gamma function. + + Matches :func:`scipy.special.rgamma`, which is 0 at the poles of gamma. + + """ + x = as_tensor_variable(x) + # gamma returns nan rather than inf at the non-positive integers, so 1 / gamma(x) + # cannot pick up the 0 that the poles call for on its own. + is_pole = le(x, 0) & eq(x, floor(x)) + return switch(is_pole, 0, 1 / gamma(x)) + + +def chdtr(v, x): + """Chi-square cumulative distribution function. + + Matches :func:`scipy.special.chdtr`. + + """ + return gammainc(v / 2, x / 2) + + +def chdtrc(v, x): + """Chi-square survival function. + + Matches :func:`scipy.special.chdtrc`. + + """ + return gammaincc(v / 2, x / 2) + + +def chdtri(v, p): + """Inverse of the chi-square survival function `chdtrc`. + + Matches :func:`scipy.special.chdtri`. + + """ + return 2 * gammainccinv(v / 2, p) + + +def gdtr(a, b, x): + """Gamma cumulative distribution function. + + Matches :func:`scipy.special.gdtr`. + + """ + return gammainc(b, a * x) + + +def gdtrc(a, b, x): + """Gamma survival function. + + Matches :func:`scipy.special.gdtrc`. + + """ + return gammaincc(b, a * x) + + +def pdtr(k, m): + """Poisson cumulative distribution function. + + Matches :func:`scipy.special.pdtr`, including its truncation of ``k``. + + """ + return gammaincc(floor(k) + 1, m) + + +def pdtrc(k, m): + """Poisson survival function. + + Matches :func:`scipy.special.pdtrc`, including its truncation of ``k``. + + """ + return gammainc(floor(k) + 1, m) + + +def _sqrt_2(): + return as_tensor_variable(np.sqrt(2.0), dtype=config.floatX) + + +class Ndtr(TensorSymbolicOp): + """Op for `ndtr`.""" + + def build_inner_graph(self, x): + return [0.5 * erfc(-x / _sqrt_2())] + + +def ndtr(x): + """Cumulative distribution function of the standard normal distribution. + + Computes the element-wise area under the standard normal density from -inf to x. + + Parameters + ---------- + x : TensorLike + Input tensor + + Returns + ------- + TensorVariable + Output tensor with the normal CDF evaluated at each element + + Examples + -------- + >>> import pytensor + >>> import pytensor.tensor as pt + >>> x = pt.vector("x") + >>> f = pytensor.function([x], pt.special.ndtr(x)) + >>> f([-1, 0, 1]) + array([0.15865525, 0.5 , 0.84134475]) + + Notes + ----- + This function corresponds to SciPy's `scipy.special.ndtr` function. Prefer + `log_ndtr` over ``log(ndtr(x))``, which underflows to ``-inf`` for x below about + -37 and saturates to ``0`` for x above about 8. + + """ + return Ndtr()(x) + + +class LogNdtr(TensorSymbolicOp): + """Op for `log_ndtr`.""" + + def build_inner_graph(self, x): + sqrt_2 = _sqrt_2() + # For x < -1 the naive log(ndtr(x)) underflows to -inf once ndtr(x) drops below + # the smallest float, so use erfc(z) = exp(-z**2) * erfcx(z), which keeps the + # tail in log space. Above it log1p avoids the cancellation as ndtr(x) -> 1. + return [ + switch( + lt(x, -1.0), + log(erfcx(-x / sqrt_2) / 2.0) - sqr(x) / 2.0, + log1p(-erfc(x / sqrt_2) / 2.0), + ) + ] + + +def log_ndtr(x): + """Logarithm of the cumulative distribution function of the standard normal. + + Stays accurate in both tails, where the naive ``log(ndtr(x))`` does not. + + Parameters + ---------- + x : TensorLike + Input tensor + + Returns + ------- + TensorVariable + Output tensor with the log normal CDF evaluated at each element + + Examples + -------- + >>> import pytensor + >>> import pytensor.tensor as pt + >>> x = pt.vector("x") + >>> f = pytensor.function([x], pt.special.log_ndtr(x)) + >>> f([-1, 0, 1]) + array([-1.84102165, -0.69314718, -0.17275378]) + + Where the naive form has lost the value entirely, this one has not: + + >>> f([-300.0, 10.0]) + array([-4.50066227e+04, -7.61985302e-24]) + + Notes + ----- + This function corresponds to SciPy's `scipy.special.log_ndtr` function. + + """ + return LogNdtr()(x) + + +def ndtri(p): + """Quantile function of the standard normal distribution. + + Inverse of :func:`ndtr`: returns the x for which ``ndtr(x) == p``. + + Parameters + ---------- + p : TensorLike + Input tensor, with values in [0, 1] + + Returns + ------- + TensorVariable + Output tensor with the normal quantile evaluated at each element + + Examples + -------- + >>> import pytensor + >>> import pytensor.tensor as pt + >>> p = pt.vector("p") + >>> f = pytensor.function([p], pt.special.ndtri(p)) + >>> f([0.1, 0.5, 0.9]) + array([-1.28155157, 0. , 1.28155157]) + + Notes + ----- + This function corresponds to SciPy's `scipy.special.ndtri` function. + """ + p = as_tensor_variable(p) + return -_sqrt_2() * erfcinv(2 * p) + + def beta(a, b): """ Beta function. @@ -228,10 +458,22 @@ def xlog1py(x, y): __all__ = [ "beta", "betaln", + "chdtr", + "chdtrc", + "chdtri", "factorial", + "gdtr", + "gdtrc", + "log_expit", + "log_ndtr", "log_softmax", "logit", + "ndtr", + "ndtri", + "pdtr", + "pdtrc", "poch", + "rgamma", "softmax", "xlog1py", "xlogy", diff --git a/tests/tensor/rewriting/test_math.py b/tests/tensor/rewriting/test_math.py index 67dc79e99f..0979675071 100644 --- a/tests/tensor/rewriting/test_math.py +++ b/tests/tensor/rewriting/test_math.py @@ -5,6 +5,7 @@ import numpy as np import pytest +import scipy.special import pytensor import pytensor.scalar as ps @@ -41,6 +42,7 @@ All, Any, Dot, + LogErfc, Max, Min, Prod, @@ -111,6 +113,8 @@ local_div_switch_sink, local_grad_log_erfc_neg, local_greedy_distributor, + local_log_erfc, + local_log_mul_erfc, local_mul_canonizer, local_mul_switch_sink, local_neg_to_mul, @@ -2565,42 +2569,56 @@ def test_local_erf_neg_minus_one(self): assert [n.op for n in f.maker.fgraph.toposort()] == [erf] def test_local_log_erfc(self): - val = [-30, -27, -26, -11, -10, -3, -2, -1, 0, 1, 2, 3, 10, 11, 26, 27, 28, 30] - if config.mode in ["DebugMode", "DEBUG_MODE", "FAST_COMPILE"]: - # python mode doesn't like the reciprocal(0) - val.remove(0) - val = np.asarray(val, dtype=config.floatX) + """``log(erfc(x))`` folds into `LogErfc`.""" x = vector("x") - # their are some `nan`s that will appear in the graph due to the logs - # of negatives values - mode = copy.copy(self.mode) - mode.check_isfinite = False - mode_fusion = copy.copy(self.mode_fusion) - mode_fusion.check_isfinite = False + result = RewriteTester( + [x], [log(erfc(x))], include=None, custom_rewrite=local_log_erfc + ) + result.assert_graph(LogErfc()(x)) + # Equivalent to the naive form wherever that one doesn't underflow + result.assert_eval(np.asarray([-3.0, -1.0, 0.0, 1.0, 3.0], dtype=config.floatX)) - f = function([x], log(erfc(x)), mode=mode) - assert len(f.maker.fgraph.apply_nodes) == 22 - assert f.maker.fgraph.outputs[0].dtype == config.floatX - assert all(np.isfinite(f(val))) + def test_local_log_erfc_scaled(self): + """``log(y * erfc(x))`` folds too, as log(y) + LogErfc(x). - f = function([x], log(erfc(-x)), mode=mode) - assert len(f.maker.fgraph.apply_nodes) == 23 - assert f.maker.fgraph.outputs[0].dtype == config.floatX - assert all(np.isfinite(f(-val))) + erfc is strictly positive, so log(y * erfc(x)) == log(y) + log(erfc(x)) holds + for any positive y, whether or not it is constant. This is the form a log of a + scaled erfc arrives in, e.g. ``ndtr(x) == 0.5 * erfc(-x / sqrt(2))``. + """ + x, y = vectors("x", "y") - f = function([x], log(erfc(x)), mode=mode_fusion) - assert len(f.maker.fgraph.apply_nodes) == 1 - assert f.maker.fgraph.outputs[0].dtype == config.floatX - assert len(f.maker.fgraph.toposort()[0].op.scalar_op.fgraph.apply_nodes) == 22 + result = RewriteTester( + [x, y], [log(y * erfc(x))], include=None, custom_rewrite=local_log_mul_erfc + ) + result.assert_graph(log(y) + LogErfc()(x)) + # Splitting into log(y) + log(erfc(x)) subtracts two near-equal numbers when + # y * erfc(x) ~ 1, so compare where the result is well conditioned. + result.assert_eval( + np.asarray([0.0, 0.5, 1.0, 2.0, 3.0], dtype=config.floatX), + np.full((5,), 0.5, dtype=config.floatX), + ) - # TODO: fix this problem: The python code upcast somewhere internally - # some value of float32 to python float for part of its computation. - # That makes the c and python code generate slightly different values - if not ( - config.floatX == "float32" and config.mode in ["DebugMode", "DEBUG_MODE"] - ): - assert all(np.isfinite(f(val))) + def test_local_log_erfc_tail(self): + """`LogErfc` must stay finite where the naive ``log(erfc(x))`` gives ``-inf``.""" + x = vector("x") + val = np.asarray([26.0, 30.0, 100.0], dtype=config.floatX) + + naive = function([x], log(erfc(x)), mode="FAST_COMPILE")(val) + assert np.isneginf(naive[1:]).all() + + mode = copy.copy(self.mode) + mode.check_isfinite = False + actual = function([x], log(erfc(x)), mode=mode)(val) + assert np.isfinite(actual).all() + assert actual.dtype == config.floatX + # erfc(x) == 2 * ndtr(-x * sqrt(2)) + expected = np.log(2) + scipy.special.log_ndtr( + -val.astype("float64") * np.sqrt(2) + ) + np.testing.assert_allclose( + actual, expected, rtol=1e-7 if config.floatX == "float64" else 1e-4 + ) @np.errstate(divide="ignore", invalid="ignore") def test_local_grad_log_erfc_neg(self): diff --git a/tests/tensor/rewriting/test_special.py b/tests/tensor/rewriting/test_special.py index 63576bb786..7957b903c2 100644 --- a/tests/tensor/rewriting/test_special.py +++ b/tests/tensor/rewriting/test_special.py @@ -11,9 +11,17 @@ from pytensor.graph.rewriting.basic import check_stack_trace from pytensor.graph.rewriting.db import RewriteDatabaseQuery from pytensor.tensor.math import exp, log -from pytensor.tensor.special import LogSoftmax, Softmax, softmax -from pytensor.tensor.type import matrix +from pytensor.tensor.rewriting.special import local_log_ndtr +from pytensor.tensor.special import ( + LogSoftmax, + Softmax, + log_ndtr, + ndtr, + softmax, +) +from pytensor.tensor.type import matrix, vector from tests import unittest_tools as utt +from tests.unittest_tools import RewriteTester _fast_run_rewrites = RewriteDatabaseQuery(include=["fast_run"]) @@ -114,3 +122,44 @@ def f(inputs): return pytensor.grad(None, x, known_grads={y: inputs}) utt.verify_grad(f, [rng.random((3, 4))]) + + +def test_local_log_ndtr(): + """``log(ndtr(x))`` folds into `LogNdtr`, the stable form.""" + x = vector("x") + + result = RewriteTester( + [x], [log(ndtr(x))], include=None, custom_rewrite=local_log_ndtr + ) + result.assert_graph(log_ndtr(x)) + # Equivalent to the naive form wherever that one doesn't lose the value + result.assert_eval( + np.asarray([-3.0, -1.0, 0.0, 1.0], dtype=config.floatX), + rtol=1e-7 if config.floatX == "float64" else 1e-5, + ) + + +def test_local_log_ndtr_tail(): + """The rewritten graph holds both tails, where the naive one loses the value. + + Below the lower tail ndtr(x) underflows to 0 and the naive form gives ``-inf``; + above the upper one it saturates to 1 and the naive form gives exactly ``0``. + """ + x = vector("x") + x_test = np.array([-300.0, -50.0, 10.0], dtype=config.floatX) + + naive = pytensor.function([x], log(ndtr(x)), mode="FAST_COMPILE")(x_test) + assert np.isneginf(naive[:2]).all() + assert naive[-1] == 0.0 + + rewritten = pytensor.function([x], log(ndtr(x)))(x_test) + assert np.isfinite(rewritten).all() + # Bit-for-bit the explicit stable form, not an approximation of it + np.testing.assert_array_equal( + rewritten, pytensor.function([x], log_ndtr(x))(x_test) + ) + np.testing.assert_allclose( + rewritten, + scipy.special.log_ndtr(x_test.astype("float64")), + rtol=1e-7 if config.floatX == "float64" else 1e-5, + ) diff --git a/tests/tensor/test_special.py b/tests/tensor/test_special.py index a19c37e5f8..a46cf074f1 100644 --- a/tests/tensor/test_special.py +++ b/tests/tensor/test_special.py @@ -1,9 +1,13 @@ import numpy as np import pytest +import scipy.special from scipy.special import beta as scipy_beta from scipy.special import factorial as scipy_factorial +from scipy.special import log_ndtr as scipy_log_ndtr from scipy.special import log_softmax as scipy_log_softmax from scipy.special import logit as scipy_logit +from scipy.special import ndtr as scipy_ndtr +from scipy.special import ndtri as scipy_ndtri from scipy.special import poch as scipy_poch from scipy.special import softmax as scipy_softmax from scipy.special import xlog1py as scipy_xlog1py @@ -13,15 +17,28 @@ from pytensor.compile.maker import function from pytensor.configdefaults import config from pytensor.graph.replace import vectorize_graph +from pytensor.tensor.math import i0e, i1e, k0, k0e, k1, k1e, log from pytensor.tensor.special import ( LogSoftmax, Softmax, beta, betaln, + chdtr, + chdtrc, + chdtri, factorial, + gdtr, + gdtrc, + log_expit, + log_ndtr, log_softmax, logit, + ndtr, + ndtri, + pdtr, + pdtrc, poch, + rgamma, softmax, xlog1py, xlogy, @@ -198,6 +215,76 @@ def test_logit(): ) +@pytest.mark.parametrize( + "pt_fn, scipy_fn, x_test", + [ + (ndtr, scipy_ndtr, np.linspace(-10, 10, 101)), + (ndtri, scipy_ndtri, np.linspace(0.01, 0.99, 101)), + (log_ndtr, scipy_log_ndtr, np.linspace(-10, 10, 101)), + ], + ids=["ndtr", "ndtri", "log_ndtr"], +) +def test_ndtr_family(pt_fn, scipy_fn, x_test): + x = vector("x") + actual_fn = function([x], pt_fn(x), allow_input_downcast=True) + np.testing.assert_allclose( + actual_fn(x_test), + scipy_fn(x_test), + rtol=1e-7 if config.floatX == "float64" else 1e-5, + ) + + +@pytest.mark.skipif( + config.floatX == "float32", reason="Tail underflows in float32 precision" +) +def test_log_ndtr_extreme_tail(): + """``log_ndtr`` must stay finite where the naive ``log(ndtr(x))`` underflows. + + Both tails are checked: below the lower one ``ndtr(x)`` underflows to 0 so the naive + form gives ``-inf``, and above the upper one it saturates to 1 so the naive form + gives exactly ``0`` and loses the value entirely. + """ + x = vector("x") + x_test = np.array([-300.0, -100.0, -50.0, -10.0, 0.0, 10.0]) + + naive = function([x], log(ndtr(x)), mode="FAST_COMPILE")(x_test) + assert np.isneginf(naive[:3]).all() + assert naive[-1] == 0.0 + + actual = function([x], log_ndtr(x))(x_test) + np.testing.assert_allclose(actual, scipy_log_ndtr(x_test), rtol=1e-7) + assert np.isfinite(actual).all() + + naive_grad = function([x], grad(log(ndtr(x)).sum(), x), mode="FAST_COMPILE")(x_test) + assert np.isnan(naive_grad[:3]).all() + + actual_grad = function([x], grad(log_ndtr(x).sum(), x))(x_test) + assert np.isfinite(actual_grad).all() + + +@pytest.mark.parametrize("pt_fn", [ndtr, log_ndtr], ids=["ndtr", "log_ndtr"]) +def test_ndtr_family_grad(pt_fn): + rng = np.random.default_rng(2321) + utt.verify_grad(pt_fn, [rng.normal(size=(3, 4))]) + + +def test_ndtri_grad(): + rng = np.random.default_rng(2321) + utt.verify_grad(ndtri, [rng.uniform(0.1, 0.9, size=(3, 4))]) + + +def test_ndtri_inverts_ndtr(): + x = vector("x") + # ndtr saturates towards 1 in the upper tail, so the roundtrip is only well + # conditioned while ndtr(x) still holds enough significant bits: at x=5, + # ndtr(x) rounds to 0.9999997, which leaves float32 nothing left to invert. + lim, tol = (5.0, 1e-7) if config.floatX == "float64" else (3.0, 1e-5) + x_test = np.linspace(-lim, lim, 51).astype(config.floatX) + np.testing.assert_allclose( + function([x], ndtri(ndtr(x)))(x_test), x_test, rtol=tol, atol=tol + ) + + def test_beta(): _a, _b = vectors("a", "b") actual_fn = function([_a, _b], beta(_a, _b)) @@ -325,3 +412,81 @@ def test_xlog1py_no_distribute_at_boundary(): f = function([a, y], xlog1py(a / 2 - 1, y)) out = f(np.int64(3), np.array([-1.0, 0.0, 1.0])) np.testing.assert_array_equal(out, np.array([-np.inf, 0.0, 0.5 * np.log(2.0)])) + + +@pytest.mark.parametrize( + "pt_fn, scipy_fn, test_values", + [ + (log_expit, scipy.special.log_expit, [np.linspace(-30, 30, 21)]), + (i0e, scipy.special.i0e, [np.linspace(-10, 10, 21)]), + (i1e, scipy.special.i1e, [np.linspace(-10, 10, 21)]), + (k0, scipy.special.k0, [np.linspace(0.1, 10, 21)]), + (k1, scipy.special.k1, [np.linspace(0.1, 10, 21)]), + (k0e, scipy.special.k0e, [np.linspace(0.1, 10, 21)]), + (k1e, scipy.special.k1e, [np.linspace(0.1, 10, 21)]), + (chdtr, scipy.special.chdtr, [np.linspace(0.5, 8, 9), np.linspace(0.1, 9, 9)]), + ( + chdtrc, + scipy.special.chdtrc, + [np.linspace(0.5, 8, 9), np.linspace(0.1, 9, 9)], + ), + ( + chdtri, + scipy.special.chdtri, + [np.linspace(0.5, 8, 9), np.linspace(0.05, 0.95, 9)], + ), + ( + gdtr, + scipy.special.gdtr, + [np.full(9, 1.5), np.linspace(0.5, 8, 9), np.linspace(0.1, 9, 9)], + ), + ( + gdtrc, + scipy.special.gdtrc, + [np.full(9, 1.5), np.linspace(0.5, 8, 9), np.linspace(0.1, 9, 9)], + ), + # non-integer k exercises the truncation scipy applies + ( + pdtr, + scipy.special.pdtr, + [np.array([0.0, 1.0, 3.0, 3.7, 9.0]), np.array([0.5, 2.0, 2.0, 2.0, 4.0])], + ), + ( + pdtrc, + scipy.special.pdtrc, + [np.array([0.0, 1.0, 3.0, 3.7, 9.0]), np.array([0.5, 2.0, 2.0, 2.0, 4.0])], + ), + # includes the poles of gamma, where rgamma is 0 but 1 / gamma(x) is nan + ( + rgamma, + scipy.special.rgamma, + [np.array([-3.0, -2.0, -1.0, 0.0, 0.5, 1.0, 4.5, 20.0])], + ), + ], + ids=[ + "log_expit", + "i0e", + "i1e", + "k0", + "k1", + "k0e", + "k1e", + "chdtr", + "chdtrc", + "chdtri", + "gdtr", + "gdtrc", + "pdtr", + "pdtrc", + "rgamma", + ], +) +def test_scipy_aliases(pt_fn, scipy_fn, test_values): + """Each alias must agree with the scipy function it claims to match.""" + inputs = [vector(f"x{i}") for i in range(len(test_values))] + fn = function(inputs, pt_fn(*inputs), allow_input_downcast=True) + np.testing.assert_allclose( + fn(*test_values), + scipy_fn(*test_values), + rtol=1e-7 if config.floatX == "float64" else 1e-5, + )