From 7d79db192581e1bf4cc239bcf3f42f80498c2f4f Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Wed, 29 Jul 2026 12:47:17 +0200 Subject: [PATCH 1/3] Only exclude symbolic_op_recognition inside SymbolicOp inner graphs Baking an OpFromGraph inner graph disabled recognition rewrites wholesale, to stop a SymbolicOp such as Softmax from re-creating itself inside its own body. That is too coarse: it also stripped recognition from ordinary OpFromGraphs, so a naive exp(x) / sum(exp(x)) that evaluates correctly at the top level returned nan once wrapped in one. Only SymbolicOps can recurse this way, so exclude recognition just for them. --- pytensor/compile/rewriting.py | 18 +++++++++--------- tests/compile/test_builders.py | 15 +++++++++++++++ 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/pytensor/compile/rewriting.py b/pytensor/compile/rewriting.py index 7c251d6b3a..7d75d6bed4 100644 --- a/pytensor/compile/rewriting.py +++ b/pytensor/compile/rewriting.py @@ -7,7 +7,7 @@ add_supervisor_to_fgraph, insert_deepcopy, ) -from pytensor.compile.builders import OpFromGraph +from pytensor.compile.builders import OpFromGraph, SymbolicOp from pytensor.compile.io import In, Out from pytensor.compile.mode import optdb from pytensor.graph.basic import Apply, Variable @@ -134,12 +134,12 @@ def rewrite_ofg_inner_graph(linker, op, node, inner, *, mode): ) -def _ofg_inner_optimizer(mode): - # Recognition rewrites fold a pattern into an inner-graph op (e.g. - # ``exp(x) / sum(exp(x))`` -> ``Softmax``, itself an ``OpFromGraph``). Running - # them on an ``OpFromGraph`` inner graph -- which may *be* that pattern -- - # would re-create the op inside itself and recurse without end. - return mode.excluding("symbolic_op_recognition").optimizer +def _ofg_inner_optimizer(mode, op): + # Recognition would re-create a `SymbolicOp` inside its own inner graph and never + # terminate. Any other `OpFromGraph` must keep it, or wrapping destabilizes the body. + if isinstance(op, SymbolicOp): + return mode.excluding("symbolic_op_recognition").optimizer + return mode.optimizer @rewrite_ofg_inner_graph.register(VMLinker) @@ -152,7 +152,7 @@ def destructive_rewrite_ofg_inner_graph(linker, op, node, inner, *, mode): # still be baked between purely internal buffers. input_specs = [In(x, borrow=True, mutable=False) for x in inner.inputs] add_supervisor_to_fgraph(fgraph=inner, input_specs=input_specs, accept_inplace=True) - _ofg_inner_optimizer(mode).rewrite(inner) + _ofg_inner_optimizer(mode, op).rewrite(inner) # The op's outputs must not alias its inputs or each other (it declares no # view_map, so the outer graph cannot see such aliases); deepcopies break any # boundary alias the optimized graph ends up with. @@ -165,7 +165,7 @@ def destructive_rewrite_ofg_inner_graph(linker, op, node, inner, *, mode): @rewrite_ofg_inner_graph.register(MLXLinker) def functional_rewrite_ofg_inner_graph(linker, op, node, inner, *, mode): """Structurally optimize the inner graph for the functional JIT backends.""" - _ofg_inner_optimizer(mode).rewrite(inner) + _ofg_inner_optimizer(mode, op).rewrite(inner) @graph_rewriter diff --git a/tests/compile/test_builders.py b/tests/compile/test_builders.py index bff5c528a2..cffdeba180 100644 --- a/tests/compile/test_builders.py +++ b/tests/compile/test_builders.py @@ -2,6 +2,7 @@ import numpy as np import pytest +import scipy.special import pytensor.tensor as pt from pytensor import Mode @@ -841,3 +842,17 @@ def test_perform_with_inner_scan_rvs(): out = op(shared(np.random.default_rng(0))) fn = function([], out, mode=Mode(linker="py", optimizer=None)) assert fn().shape == (2,) + + +def test_inner_graph_keeps_symbolic_op_recognition(): + # Wrapping an expression in an OpFromGraph must not destabilize it: the inner graph + # still needs the rewrite recognizing exp(x) / sum(exp(x)) as the stable Softmax. + x = pt.vector("x") + inner = pt.vector("inner") + e = exp(inner) + op = OpFromGraph([inner], [e / e.sum()]) + + test_val = np.array([1000.0, 1001.0]) + np.testing.assert_allclose( + function([x], op(x))(test_val), scipy.special.softmax(test_val) + ) From 3ba606b3b0448011675807fd75e0fd32e14126a2 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Wed, 29 Jul 2026 12:47:06 +0200 Subject: [PATCH 2/3] Make logsumexp and logaddexp stable by construction pt.logsumexp built a bare log(sum(exp(x))) and left the max subtraction to local_log_sum_exp. That rewrite runs after grad, so the gradient was built from the raw form as exp(x) / sum(exp(x)), which overflows past log(float64 max) and is 0/0 for very negative x. Issue #2288 only addressed the forward value. Add a LogSumExp SymbolicOp that subtracts the maximum as built, and route logsumexp and logaddexp through it. Differentiating the stabilized body is already stable, so no pullback override is needed. Inlined at specialize by late_inline_OpFromGraph, as for XLogY. local_log_sum_exp and local_log_add_exp now emit that Op rather than expanding the stable form inline, so hand-written and helper-built forms share one representation. They are tagged symbolic_op_recognition, which keeps them from firing inside LogSumExp's own body, and registered at stabilize only: at specialize they would re-wrap what late_inline_OpFromGraph had just inlined. --- pytensor/tensor/math.py | 14 +++++++-- pytensor/tensor/rewriting/math.py | 52 +++++-------------------------- pytensor/tensor/rewriting/ofg.py | 4 +-- pytensor/tensor/special.py | 27 ++++++++++++++++ tests/tensor/test_math.py | 18 +++++++++++ 5 files changed, 67 insertions(+), 48 deletions(-) diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py index daeb279251..5a3c5a2586 100644 --- a/pytensor/tensor/math.py +++ b/pytensor/tensor/math.py @@ -3906,8 +3906,10 @@ def logaddexp(*xs): TensorVariable """ + from pytensor.tensor.extra_ops import broadcast_arrays - return log(add(*[exp(x) for x in xs])) + xs = broadcast_arrays(*[as_tensor_variable(x) for x in xs]) + return logsumexp(stack(xs, axis=0), axis=0) def logsumexp(x, axis=None, keepdims=False): @@ -3934,8 +3936,16 @@ def logsumexp(x, axis=None, keepdims=False): TensorVariable """ + from pytensor.tensor.special import LogSumExp - return log(sum(exp(x), axis=axis, keepdims=keepdims)) + x = as_tensor_variable(x) + axis = ( + tuple(range(x.type.ndim)) + if axis is None + else normalize_axis_tuple(axis, x.type.ndim) + ) + out = LogSumExp(axis=axis)(x) + return expand_dims(out, axis) if keepdims else out _matmul = Blockwise(_dot, name="Matmul") diff --git a/pytensor/tensor/rewriting/math.py b/pytensor/tensor/rewriting/math.py index 9edee1368b..fa05a9ad25 100644 --- a/pytensor/tensor/rewriting/math.py +++ b/pytensor/tensor/rewriting/math.py @@ -71,7 +71,6 @@ expm1, ge, int_div, - isinf, ive, kve, le, @@ -79,8 +78,8 @@ log1mexp, log1p, log1pexp, - makeKeepDims, - maximum, + logaddexp, + logsumexp, mul, neg, polygamma, @@ -100,9 +99,7 @@ variadic_mul, ) from pytensor.tensor.math import abs as pt_abs -from pytensor.tensor.math import max as pt_max from pytensor.tensor.math import pow as pt_pow -from pytensor.tensor.math import sum as pt_sum from pytensor.tensor.rewriting.basic import ( broadcast_like_elemwise, local_second_sink, @@ -2813,43 +2810,26 @@ def local_log1p(fgraph, node): return [broadcast_like_elemwise(new_out, node, fgraph=fgraph, stack_trace=True)] -@register_stabilize("fast_compile") -@register_specialize +@register_stabilize("symbolic_op_recognition", "fast_compile") @node_rewriter([log]) def local_log_add_exp(fgraph, node): - """ - ``log(exp(x)+exp(y)+exp(z)) = max + log(x-max, y-max, z-max)`` + """``log(exp(x) + exp(y) + exp(z)) -> logaddexp(x, y, z)``. 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): - # Do not offset when max_pre = -np.inf, to avoid nan in the output - # Switch statement is placed directly inside add to break the self-symmetry - # of the returned output (otherwise the rewrite would not stabilize) - max_pre = reduce(maximum, pre_exp) - ret = max_pre + log( - add( - *[ - switch(isinf(max_pre), exp(max_pre), exp(p - max_pre)) - for p in pre_exp - ] - ) - ) - return [ret] + return [logaddexp(*pre_exp)] -@register_stabilize("fast_compile") -@register_specialize +@register_stabilize("symbolic_op_recognition", "fast_compile") @node_rewriter([log]) def local_log_sum_exp(fgraph, node): - # log(sum_i(exp(x_i))) = x_max + log(sum_i(exp(x_i - x_max))) - + """``log(sum_i(exp(x_i))) -> logsumexp(x)``.""" sum_node = node.inputs[0].owner # If the sum has keepdims=True, there might be a dimshuffle if sum_node and isinstance(sum_node.op, DimShuffle): @@ -2869,23 +2849,7 @@ def local_log_sum_exp(fgraph, node): ): return - pre_exp = exp_node.inputs[0] - max_pre_exp = pt_max(pre_exp, axis=axis) - max_pre_exp_keepdims = makeKeepDims(pre_exp, max_pre_exp, axis) - - # Do not offset when max_pre = -np.inf, to avoid nan in the output - # Switch statement is placed directly inside sum to break the self-symmetry - # of the returned output (otherwise the rewrite would not stabilize) - ret = max_pre_exp + log( - pt_sum( - switch( - isinf(max_pre_exp_keepdims), - exp(max_pre_exp_keepdims), - exp(pre_exp - max_pre_exp_keepdims), - ), - axis=axis, - ), - ) + ret = logsumexp(exp_node.inputs[0], axis=axis) # Restore the dimshuffle op, if any. if dimshuffle_op: diff --git a/pytensor/tensor/rewriting/ofg.py b/pytensor/tensor/rewriting/ofg.py index 5f8b865fee..b68e842377 100644 --- a/pytensor/tensor/rewriting/ofg.py +++ b/pytensor/tensor/rewriting/ofg.py @@ -2,11 +2,11 @@ from pytensor.graph import node_rewriter from pytensor.tensor.basic import AllocDiag from pytensor.tensor.rewriting.basic import register_specialize -from pytensor.tensor.special import XLog1PY, XLogY +from pytensor.tensor.special import LogSumExp, XLog1PY, XLogY @register_specialize("inline_ofg") -@node_rewriter([AllocDiag, XLogY, XLog1PY]) +@node_rewriter([AllocDiag, XLogY, XLog1PY, LogSumExp]) def late_inline_OpFromGraph(fgraph, node): """ Inline `OpFromGraph` nodes. diff --git a/pytensor/tensor/special.py b/pytensor/tensor/special.py index 2cc1acbf52..711a95afab 100644 --- a/pytensor/tensor/special.py +++ b/pytensor/tensor/special.py @@ -9,6 +9,7 @@ exp, gamma, gammaln, + isinf, log, log1p, mul, @@ -96,6 +97,32 @@ def log_softmax(c, axis=None): return LogSoftmax(axis=axis)(c) +class LogSumExp(TensorSymbolicOp): + r"""Log of the sum of exponentials. + + :math:`\log \sum_k e^{x_k}` + + Includes the numerical stabilization trick (subtracting the maximum), so unlike a + bare ``log(sum(exp(x)))`` the gradient taken from it is stable too. + """ + + __props__ = ("axis",) + + # See the note on `XLogY.inline`. + inline = False + + def __init__(self, *, axis, **kwargs): + self.axis = tuple(axis) + super().__init__(**kwargs) + + def build_inner_graph(self, x): + x_max = x.max(axis=self.axis, keepdims=True) + # Do not offset when x_max = -inf, to avoid nan in the output + x_max = switch(isinf(x_max), 0.0, x_max) + out = log(exp(x - x_max).sum(axis=self.axis, keepdims=True)) + x_max + return [out.squeeze(axis=self.axis)] + + @_vectorize_node.register(Softmax) @_vectorize_node.register(LogSoftmax) def vectorize_softmax_node(op, node, batched_x): diff --git a/tests/tensor/test_math.py b/tests/tensor/test_math.py index 8ade530891..62319bb9b5 100644 --- a/tests/tensor/test_math.py +++ b/tests/tensor/test_math.py @@ -3771,6 +3771,24 @@ def test_logsumexp(shape, axis, keepdims): ) +@pytest.mark.parametrize("mode", ["FAST_RUN", "FAST_COMPILE"]) +def test_logsumexp_logaddexp_stable_grad(mode): + """Both helpers promise a stable forward, so their gradient must be stable too. + + A naive ``log(sum(exp(x)))`` is only stabilized by `local_log_sum_exp`, which runs + too late to help `grad`: the gradient built from the raw form is + ``exp(x) / sum(exp(x))``, which overflows for large ``x`` and is 0/0 for very + negative ``x``. Multiplying by ``w`` keeps the output gradient from folding to 1, + which is what would otherwise let `local_softmax_stabilize` repair the pattern. + """ + x, w = vector("x"), scalar("w") + + for out in (logsumexp(x), logaddexp(x[0], x[1])): + f = function([x, w], grad(out * w, x), mode=mode) + for inp in (np.array([1000.0, 1001.0]), np.array([-800.0, -805.0])): + np.testing.assert_allclose(f(inp, 2.0), 2.0 * scipy.special.softmax(inp)) + + def test_pprint(): x = vector("x") y = pt_sum(x, axis=0) From 0b9a7737895062bcb75a9aa300c48df8def17447 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Wed, 29 Jul 2026 12:47:18 +0200 Subject: [PATCH 3/3] Recognize log_softmax and rewrite exp(log_softmax) -> softmax --- pytensor/graph/features.py | 8 ++--- pytensor/graph/rewriting/unify.py | 3 +- pytensor/tensor/math.py | 6 +--- pytensor/tensor/rewriting/special.py | 40 ++++++++++++++++++--- pytensor/tensor/special.py | 9 ++--- pytensor/tensor/utils.py | 12 +++++-- tests/tensor/rewriting/test_special.py | 49 ++++++++++++++++++++++++-- 7 files changed, 102 insertions(+), 25 deletions(-) diff --git a/pytensor/graph/features.py b/pytensor/graph/features.py index 47b90dea02..43486bdc35 100644 --- a/pytensor/graph/features.py +++ b/pytensor/graph/features.py @@ -546,10 +546,10 @@ class FullHistory(Feature): └─ ··· >> local_softmax_stabilize Log [id A] 1 - └─ Softmax{axis=None} [id B] 0 + └─ Softmax{axis=(0,)} [id B] 0 └─ x [id C] >> local_logsoftmax - LogSoftmax{axis=None} [id A] 0 + LogSoftmax{axis=(0,)} [id A] 0 └─ x [id B] @@ -563,7 +563,7 @@ class FullHistory(Feature): .. testoutput:: >> local_logsoftmax Log [id A] 1 - └─ Softmax{axis=None} [id B] 0 + └─ Softmax{axis=(0,)} [id B] 0 └─ x [id C] >> local_softmax_stabilize Log [id A] 4 @@ -591,7 +591,7 @@ class FullHistory(Feature): .. testoutput:: Log [id A] 1 - └─ Softmax{axis=None} [id B] 0 + └─ Softmax{axis=(0,)} [id B] 0 └─ x [id C] diff --git a/pytensor/graph/rewriting/unify.py b/pytensor/graph/rewriting/unify.py index 26572c2346..27a25d4155 100644 --- a/pytensor/graph/rewriting/unify.py +++ b/pytensor/graph/rewriting/unify.py @@ -572,7 +572,8 @@ def reify_pattern(pattern, subs: Mapping[PatternVar | Asterisk, Any]): inputs.extend(captured) else: inputs.append(reify_pattern(p, subs)) - return op.make_node(*inputs).default_output() + # Call the Op, so those that build state lazily (SymbolicOp) are constructed + return op(*inputs) if isinstance(pattern, OpPattern): op_type = pattern.op_type diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py index 5a3c5a2586..9a4a1ec92f 100644 --- a/pytensor/tensor/math.py +++ b/pytensor/tensor/math.py @@ -3939,11 +3939,7 @@ def logsumexp(x, axis=None, keepdims=False): from pytensor.tensor.special import LogSumExp x = as_tensor_variable(x) - axis = ( - tuple(range(x.type.ndim)) - if axis is None - else normalize_axis_tuple(axis, x.type.ndim) - ) + axis = normalize_reduce_axis(axis, x.type.ndim, normalize_none=True) out = LogSumExp(axis=axis)(x) return expand_dims(out, axis) if keepdims else out diff --git a/pytensor/tensor/rewriting/special.py b/pytensor/tensor/rewriting/special.py index 067b839851..6be03c77f4 100644 --- a/pytensor/tensor/rewriting/special.py +++ b/pytensor/tensor/rewriting/special.py @@ -1,14 +1,20 @@ -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.math import Sum, exp, log, sub, true_div from pytensor.tensor.rewriting.basic import register_stabilize -from pytensor.tensor.special import Softmax, log_softmax +from pytensor.tensor.special import LogSoftmax, LogSumExp, Softmax, log_softmax from pytensor.tensor.subtensor import ( AdvancedSubtensor, Subtensor, ) from pytensor.tensor.type import values_eq_approx_remove_inf +from pytensor.tensor.utils import normalize_reduce_axis subtensor_ops = ( @@ -64,6 +70,32 @@ def find_softmax_under_lifteable_ops(inp_node, ops_to_lift): return [ret] +# Exp(LogSoftmax(x)) -> Softmax(x) +local_exp_log_softmax = PatternNodeRewriter( + (exp, (OpPattern(LogSoftmax, axis="axis"), "x")), + (OpPattern(Softmax, axis="axis"), "x"), + name="local_exp_log_softmax", +) +register_stabilize(local_exp_log_softmax) + + +# x - logsumexp(x, axis, keepdims=True) -> LogSoftmax(x) +# The shared "axis" makes the DimShuffle match only when it re-expands the reduced axes +local_log_softmax_from_logsumexp = PatternNodeRewriter( + ( + sub, + "x", + ( + OpPattern(DimShuffle, is_expand_dims=True, augment="axis"), + (OpPattern(LogSumExp, axis="axis"), "x"), + ), + ), + (OpPattern(LogSoftmax, axis="axis"), "x"), + name="local_log_softmax_from_logsumexp", +) +register_stabilize(local_log_softmax_from_logsumexp) + + @register_stabilize("symbolic_op_recognition", "fast_compile") @node_rewriter([true_div]) def local_softmax_stabilize(fgraph, node): @@ -92,6 +124,6 @@ def local_softmax_stabilize(fgraph, node): case _: return None - ret = Softmax(axis=axis)(x) + ret = Softmax(axis=normalize_reduce_axis(axis, x.type.ndim, normalize_none=True))(x) copy_stack_trace(node.outputs, ret) return [ret] diff --git a/pytensor/tensor/special.py b/pytensor/tensor/special.py index 711a95afab..dd43ac40ae 100644 --- a/pytensor/tensor/special.py +++ b/pytensor/tensor/special.py @@ -1,5 +1,3 @@ -from numpy.lib.array_utils import normalize_axis_tuple - from pytensor.gradient import DisconnectedType, disconnected_type from pytensor.graph.replace import _vectorize_node from pytensor.tensor import as_tensor_variable @@ -17,6 +15,7 @@ switch, ) from pytensor.tensor.symbolic import TensorSymbolicOp +from pytensor.tensor.utils import normalize_reduce_axis class Softmax(TensorSymbolicOp): @@ -56,8 +55,7 @@ def pushforward(self, inputs, outputs, eval_points): def softmax(c, axis=None): c = as_tensor_variable(c) - if axis is not None: - axis = normalize_axis_tuple(axis, c.type.ndim) + axis = normalize_reduce_axis(axis, c.type.ndim, normalize_none=True) return Softmax(axis=axis)(c) @@ -92,8 +90,7 @@ def pullback(self, inputs, outputs, output_grads): def log_softmax(c, axis=None): c = as_tensor_variable(c) - if axis is not None: - axis = normalize_axis_tuple(axis, c.type.ndim) + axis = normalize_reduce_axis(axis, c.type.ndim, normalize_none=True) return LogSoftmax(axis=axis)(c) diff --git a/pytensor/tensor/utils.py b/pytensor/tensor/utils.py index 8662faded3..99531ca6d7 100644 --- a/pytensor/tensor/utils.py +++ b/pytensor/tensor/utils.py @@ -229,10 +229,16 @@ def operand_sig(operand_ndim: int, prefix: str) -> str: return f"{inputs_sig}->{outputs_sig}" -def normalize_reduce_axis(axis, ndim: int) -> tuple[int, ...] | None: - """Normalize the axis parameter for reduce operations.""" +def normalize_reduce_axis( + axis, ndim: int, normalize_none: bool = False +) -> tuple[int, ...] | None: + """Normalize the axis parameter for reduce operations. + + With ``normalize_none`` a ``None`` axis becomes every axis, for Ops that need them + spelled out so that ``None`` and the explicit axes give Ops that compare equal. + """ if axis is None: - return None + return tuple(range(ndim)) if normalize_none else None # scalar inputs are treated as 1D regarding axis in reduce operations if axis is not None: diff --git a/tests/tensor/rewriting/test_special.py b/tests/tensor/rewriting/test_special.py index 63576bb786..10651e9395 100644 --- a/tests/tensor/rewriting/test_special.py +++ b/tests/tensor/rewriting/test_special.py @@ -10,10 +10,15 @@ from pytensor.graph.fg import FunctionGraph 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.math import exp, log, logsumexp +from pytensor.tensor.rewriting.special import ( + local_exp_log_softmax, + local_log_softmax_from_logsumexp, +) +from pytensor.tensor.special import LogSoftmax, Softmax, log_softmax, softmax from pytensor.tensor.type import matrix from tests import unittest_tools as utt +from tests.unittest_tools import RewriteTester _fast_run_rewrites = RewriteDatabaseQuery(include=["fast_run"]) @@ -38,6 +43,46 @@ def test_local_logsoftmax_rewrite(self, axis): assert check_stack_trace(fgraph, ops_to_check=LogSoftmax) assert check_stack_trace(fgraph, ops_to_check="all") + @pytest.mark.parametrize("axis", [None, 0, -1]) + def test_local_exp_log_softmax_rewrite(self, axis): + """Check that ``Exp(LogSoftmax(x)) -> Softmax(x)``.""" + x = matrix("x") + test_val = np.array([[1000.0, 1001.0], [1.0, 2.0]]) + + result = RewriteTester( + [x], [exp(log_softmax(x, axis=axis))], custom_rewrite=local_exp_log_softmax + ) + result.assert_graph(softmax(x, axis=axis)) + result.assert_eval(test_val) + + # When the LogSoftmax is needed anyway, reusing it beats repeating the reduction + log_sm = log_softmax(x, axis=axis) + result = RewriteTester( + [x], [log_sm, exp(log_sm)], custom_rewrite=local_exp_log_softmax + ) + result.assert_graph(log_sm, exp(log_sm)) + + @pytest.mark.parametrize("axis", [None, 0, -1]) + def test_local_log_softmax_from_logsumexp(self, axis): + """Check that ``x - logsumexp(x, axis, keepdims=True) -> LogSoftmax(x)``.""" + x = matrix("x") + test_val = np.array([[1000.0, 1001.0], [1.0, 2.0]]) + + result = RewriteTester( + [x], + [x - logsumexp(x, axis=axis, keepdims=True)], + custom_rewrite=local_log_softmax_from_logsumexp, + ) + result.assert_graph(log_softmax(x, axis=axis)) + result.assert_eval(test_val) + + # When the logsumexp is needed anyway, reusing it beats repeating the reduction + lse = logsumexp(x, axis=axis, keepdims=True) + result = RewriteTester( + [x], [x - lse, lse], custom_rewrite=local_log_softmax_from_logsumexp + ) + result.assert_graph(x - lse, lse) + @pytest.mark.parametrize("axis", [None, 0, -1]) @pytest.mark.parametrize("idx0", [0, slice(1, None), slice(None)]) @pytest.mark.parametrize("idx1", [None, [0, 1, 1, -1]])