scipy compat

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

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"

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

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",

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

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.

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