From 90dc8e24e127b8fc2826a65f518167c1f377188d Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Sun, 9 Aug 2026 09:50:09 +0200 Subject: [PATCH] Log exp stabilizations with mixed domain inputs --- pytensor/tensor/rewriting/special.py | 125 +++++++++++++++++-- tests/tensor/rewriting/test_special.py | 160 ++++++++++++++++++++++++- 2 files changed, 271 insertions(+), 14 deletions(-) diff --git a/pytensor/tensor/rewriting/special.py b/pytensor/tensor/rewriting/special.py index 84aa1ad76c..4291173dc0 100644 --- a/pytensor/tensor/rewriting/special.py +++ b/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() - 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") diff --git a/tests/tensor/rewriting/test_special.py b/tests/tensor/rewriting/test_special.py index f2293c3e54..724cba8789 100644 --- a/tests/tensor/rewriting/test_special.py +++ b/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: