Log exp stabilizations with mixed domain inputs

pytensor/tensor/rewriting/special.py

@@ -1,14 +1,19 @@
+import numpy as np
+
+from pytensor.graph.basic import Constant
 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.scalar.basic import Add, Exp
+from pytensor.tensor.basic import cast
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
-from pytensor.tensor.math import Sum, add, exp, log, sub, true_div
-from pytensor.tensor.rewriting.basic import register_stabilize
+from pytensor.tensor.math import Sum, exp, log, sigmoid, softplus, sub, true_div
+from pytensor.tensor.rewriting.basic import register_specialize, register_stabilize
 from pytensor.tensor.special import (
+    LogAddExp,
     LogSoftmax,
     LogSumExp,
     Softmax,
@@ -19,6 +24,7 @@
 from pytensor.tensor.subtensor import (
     AdvancedSubtensor,
     Subtensor,
+    _is_provably_positive,
 )
 from pytensor.tensor.type import values_eq_approx_remove_inf
 from pytensor.tensor.utils import normalize_reduce_axis
@@ -136,20 +142,117 @@ def local_softmax_stabilize(fgraph, node):
     return [ret]
 
 
+def _addends_in_log_form(terms):
+    """Map each addend to its logarithm, or return ``None`` when that is unsafe.
+
+    Addends that are ``exp(x)`` map to ``x``; the rest must be provably positive
+    (e.g. positive constants) and map to ``log(term)``, which constant folding
+    evaluates for constants. At least one addend must be an ``exp`` for the sum
+    to be at risk of overflow, otherwise ``None`` is returned.
+    """
+    log_terms = []
+    any_exp = False
+    for term in terms:
+        match term.owner_op_and_inputs:
+            case Elemwise(Exp()), x:
+                log_terms.append(x)
+                any_exp = True
+            case _ if _is_provably_positive(term):
+                log_terms.append(log(term))
+            case _:
+                return None
+    return log_terms if any_exp else None
+
+
 @register_stabilize("symbolic_op_recognition", "fast_compile")
 @node_rewriter([log])
 def local_log_add_exp(fgraph, node):
-    """``log(exp(x) + exp(y) + exp(z)) -> logaddexp(x, y, z)``.
+    """``log(exp(x) + exp(y) + c) -> logaddexp(x, y, log(c))``.
+
+    Every addend must be an ``exp`` or provably positive (e.g. a positive constant),
+    and at least one must be an ``exp``.
 
     TODO: in canonicalize, change log10 and log2 -> log
     """
-    z = node.inputs[0]
-    if z.owner and z.owner.op == add:
-        zi = z.owner.inputs
-        pre_exp = [x.owner.inputs[0] for x in zi if x.owner and x.owner.op == exp]
-        # all arguments to add are exp(<something>)
-        if len(pre_exp) == len(zi):
-            return [logaddexp(*pre_exp)]
+    match node.inputs[0].owner_op_and_inputs:
+        case Elemwise(Add()), *terms:
+            pass
+        case _:
+            return None
+
+    log_terms = _addends_in_log_form(terms)
+    if log_terms is None:
+        return None
+
+    ret = logaddexp(*log_terms)
+    if ret.dtype != node.outputs[0].dtype:
+        ret = cast(ret, node.outputs[0].dtype)
+    copy_stack_trace(node.outputs, ret)
+    return [ret]
+
+
+@register_stabilize("symbolic_op_recognition", "fast_compile")
+@node_rewriter([true_div])
+def local_sigmoid_stabilize(fgraph, node):
+    """Detect ``exp(x) / (exp(x) + exp(y) + c)`` and replace it with
+    ``sigmoid(x - logaddexp(y, log(c)))``.
+
+    The numerator must be one of the denominator addends and the same restrictions of
+    `local_log_add_exp` apply to the addends. With a two-addend denominator the
+    ``logaddexp`` collapses to the log form of the other addend, so
+    ``c / (c + exp(x)) -> sigmoid(log(c) - x)`` and
+    ``exp(x) / (exp(x) + exp(y)) -> sigmoid(x - y)``.
+    """
+    num, denom = node.inputs
+    match denom.owner_op_and_inputs:
+        case Elemwise(Add()), *terms:
+            pass
+        case _:
+            return None
+
+    num_idx = next((i for i, term in enumerate(terms) if term is num), None)
+    if num_idx is None:
+        return None
+
+    log_terms = _addends_in_log_form(terms)
+    if log_terms is None:
+        return None
+
+    others = [term for i, term in enumerate(log_terms) if i != num_idx]
+    ret = sigmoid(
+        log_terms[num_idx] - (logaddexp(*others) if len(others) > 1 else others[0])
+    )
+    if ret.dtype != node.outputs[0].dtype:
+        ret = cast(ret, node.outputs[0].dtype)
+    copy_stack_trace(node.outputs, ret)
+    return [ret]
+
+
+@register_stabilize
+@register_specialize
+@node_rewriter([LogAddExp])
+def local_logaddexp_const_to_softplus(fgraph, node):
+    """``logaddexp(c, x) -> c + softplus(x - c)`` for a finite constant ``c``.
+
+    A single ``softplus`` replaces the max-subtraction inner graph of `LogAddExp`, and
+    its gradient is a bare ``sigmoid``. Restricted to finite constants: with
+    ``c = -inf`` this form would give ``nan`` where `LogAddExp` correctly returns
+    ``x``.
+
+    Registered in stabilize so it sees the `LogAddExp` emitted by `local_log_add_exp`
+    (whose ``log(c)`` input end-of-pass constant folding collapses) before the
+    specialize-time `OpFromGraph` inliner gets to it.
+    """
+    if len(node.inputs) != 2:
+        return None
+    for const, x in (node.inputs, node.inputs[::-1]):
+        if isinstance(const, Constant) and np.isfinite(const.data).all():
+            ret = const + softplus(x - const)
+            if ret.dtype != node.outputs[0].dtype:
+                ret = cast(ret, node.outputs[0].dtype)
+            copy_stack_trace(node.outputs, ret)
+            return [ret]
+    return None
 
 
 @register_stabilize("symbolic_op_recognition", "fast_compile")

tests/tensor/rewriting/test_special.py

@@ -13,15 +13,28 @@
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.basic import check_stack_trace
 from pytensor.graph.rewriting.db import RewriteDatabaseQuery
+from pytensor.scalar.math import Softplus
+from pytensor.tensor.basic import constant
 from pytensor.tensor.elemwise import DimShuffle
-from pytensor.tensor.math import Max, exp, log
+from pytensor.tensor.math import Max, add, exp, log, sigmoid, softplus
 from pytensor.tensor.math import sum as pt_sum
 from pytensor.tensor.rewriting.special import (
     local_exp_log_softmax,
+    local_log_add_exp,
     local_log_softmax_from_logsumexp,
+    local_logaddexp_const_to_softplus,
+    local_sigmoid_stabilize,
 )
-from pytensor.tensor.special import LogSoftmax, Softmax, log_softmax, logsumexp, softmax
-from pytensor.tensor.type import TensorType, dvector, matrix, tensor3, vector
+from pytensor.tensor.special import (
+    LogAddExp,
+    LogSoftmax,
+    Softmax,
+    log_softmax,
+    logaddexp,
+    logsumexp,
+    softmax,
+)
+from pytensor.tensor.type import TensorType, dscalar, dvector, matrix, tensor3, vector
 from tests import unittest_tools as utt
 from tests.unittest_tools import RewriteTester
 
@@ -200,9 +213,150 @@ def test_local_log_add_exp(mode):
     assert np.isfinite(f([10000], [10000]))  # causes overflow if handled incorrectly
     utt.assert_allclose(f([10000], [10000]), 20000)
 
+    # test that provably positive non-exp addends are handled as exp(log(c))
+    x = dvector()
+    f = pytensor.function([x], log(2.0 + exp(x)), mode=m)
+
+    utt.assert_allclose(f([0]), np.log(3.0))
+    assert np.isfinite(f([800]))  # causes overflow if handled incorrectly
+    utt.assert_allclose(f([800]), 800.0)
+
     # TODO: test that the rewrite works in the presence of broadcasting.
 
 
+def test_local_log_add_exp_positive_addends():
+    """``log(c + exp(x)) -> logaddexp(log(c), x)`` for provably positive ``c``."""
+    x = dscalar("x")
+    c = constant(2.0)
+
+    result = RewriteTester(
+        [x], [log(add(c, exp(x)))], include=None, custom_rewrite=local_log_add_exp
+    )
+    result.assert_graph(logaddexp(log(c), x))
+    result.assert_eval(3.0)
+
+    # The unstable form overflows for x past log(float64 max)
+    np.testing.assert_allclose(result.rewr_fn(800.0), 800.0)
+
+    # A negative constant or a sign-unknown variable blocks the rewrite
+    y = dscalar("y")
+    for other in (constant(-2.0), y):
+        unstable = log(add(other, exp(x)))
+        result = RewriteTester(
+            [x, y], [unstable], include=None, custom_rewrite=local_log_add_exp
+        )
+        result.assert_graph(unstable)
+
+    # So does the absence of any exp addend
+    result = RewriteTester(
+        [x], [log(add(c, sigmoid(x)))], include=None, custom_rewrite=local_log_add_exp
+    )
+    result.assert_graph(log(add(c, sigmoid(x))))
+
+
+def test_local_sigmoid_stabilize():
+    """Divisions by a sum of exp / provably positive addends that includes the
+    numerator are recognized as ``sigmoid``.
+    """
+    x = dscalar("x")
+    y = dscalar("y")
+    c = constant(2.0)
+
+    # c / (c + exp(x)) -> sigmoid(log(c) - x)
+    result = RewriteTester(
+        [x], [c / add(c, exp(x))], include=None, custom_rewrite=local_sigmoid_stabilize
+    )
+    result.assert_graph(sigmoid(log(c) - x))
+    result.assert_eval(3.0)
+
+    # exp(x) / (exp(x) + c) -> sigmoid(x - log(c))
+    ex = exp(x)
+    result = RewriteTester(
+        [x], [ex / add(ex, c)], include=None, custom_rewrite=local_sigmoid_stabilize
+    )
+    result.assert_graph(sigmoid(x - log(c)))
+    result.assert_eval(3.0)
+    # The unstable form is inf / inf -> nan for x past log(float64 max)
+    np.testing.assert_allclose(result.rewr_fn(800.0), 1.0)
+
+    # exp(x) / (exp(x) + exp(y)) -> sigmoid(x - y)
+    ey = exp(y)
+    result = RewriteTester(
+        [x, y], [ex / add(ex, ey)], include=None, custom_rewrite=local_sigmoid_stabilize
+    )
+    result.assert_graph(sigmoid(x - y))
+    result.assert_eval(3.0, 2.5)
+    np.testing.assert_allclose(result.rewr_fn(800.0, 800.0), 0.5)
+
+    # Remaining addends are combined with logaddexp:
+    # exp(x) / (exp(x) + exp(y) + c) -> sigmoid(x - logaddexp(y, log(c)))
+    result = RewriteTester(
+        [x, y],
+        [ex / add(ex, ey, c)],
+        include=None,
+        custom_rewrite=local_sigmoid_stabilize,
+    )
+    result.assert_graph(sigmoid(x - logaddexp(y, log(c))))
+    result.assert_eval(3.0, 2.5)
+
+    # A numerator that is not one of the addends blocks the rewrite
+    result = RewriteTester(
+        [x, y], [ey / add(ex, c)], include=None, custom_rewrite=local_sigmoid_stabilize
+    )
+    result.assert_graph(ey / add(ex, c))
+
+    # So does a sign-unknown addend
+    result = RewriteTester(
+        [x, y], [ex / add(ex, y)], include=None, custom_rewrite=local_sigmoid_stabilize
+    )
+    result.assert_graph(ex / add(ex, y))
+
+
+def test_local_logaddexp_const_to_softplus():
+    """``logaddexp(c, x) -> c + softplus(x - c)`` for a finite constant ``c``."""
+    x = dscalar("x")
+    c = constant(2.0)
+
+    for out in (logaddexp(c, x), logaddexp(x, c)):
+        result = RewriteTester(
+            [x], [out], include=None, custom_rewrite=local_logaddexp_const_to_softplus
+        )
+        result.assert_graph(c + softplus(x - c))
+        result.assert_eval(3.0)
+        np.testing.assert_allclose(result.rewr_fn(800.0), 800.0)
+
+    # c = -inf needs LogAddExp's max-subtraction guard: the softplus form would be nan
+    neg_inf = constant(-np.inf)
+    result = RewriteTester(
+        [x],
+        [logaddexp(neg_inf, x)],
+        include=None,
+        custom_rewrite=local_logaddexp_const_to_softplus,
+    )
+    result.assert_graph(logaddexp(neg_inf, x))
+
+    # Variadic logaddexp is left alone
+    y = dscalar("y")
+    result = RewriteTester(
+        [x, y],
+        [logaddexp(c, x, y)],
+        include=None,
+        custom_rewrite=local_logaddexp_const_to_softplus,
+    )
+    result.assert_graph(logaddexp(c, x, y))
+
+    # In the full pipeline recognition plus lowering leave a plain softplus graph
+    f = pytensor.function(
+        [x], log(2.0 + exp(x)), mode=get_mode("FAST_RUN").excluding("fusion")
+    )
+    topo = f.maker.fgraph.toposort()
+    assert not any(isinstance(node.op, LogAddExp) for node in topo)
+    assert any(
+        isinstance(getattr(node.op, "scalar_op", None), Softplus) for node in topo
+    )
+    np.testing.assert_allclose(f(800.0), 800.0)
+
+
 def compile_graph_log_sum_exp(x, axis, dimshuffle_op=None, mode="FAST_RUN"):
     sum_exp = pt_sum(exp(x), axis=axis)
     if dimshuffle_op: