Stable logsumexp softmax gradients without rewrites

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

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]
 
 

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

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,12 @@ 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 = normalize_reduce_axis(axis, x.type.ndim, normalize_none=True)
+    out = LogSumExp(axis=axis)(x)
+    return expand_dims(out, axis) if keepdims else out
 
 
 _matmul = Blockwise(_dot, name="Matmul")

pytensor/tensor/rewriting/math.py

@@ -71,16 +71,15 @@
     expm1,
     ge,
     int_div,
-    isinf,
     ive,
     kve,
     le,
     log,
     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(<something>)
         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:

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.

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]

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
@@ -9,13 +7,15 @@
     exp,
     gamma,
     gammaln,
+    isinf,
     log,
     log1p,
     mul,
     sum,
     switch,
 )
 from pytensor.tensor.symbolic import TensorSymbolicOp
+from pytensor.tensor.utils import normalize_reduce_axis
 
 
 class Softmax(TensorSymbolicOp):
@@ -55,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)
 
 
@@ -91,11 +90,36 @@ 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)
 
 
+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):

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: