PERF: map pt.logsumexp pattern to native mx.logsumexp (1.5× speedup)

Part of #2085.

Current implementation

pt.logsumexp(x, axis=...) is not a single op — it's built in pytensor/tensor/math.py from Max + Exp + Sum + Log + Add, with an Isinf/Switch guard for -inf inputs. After lower_xtensor and canonicalization, the resulting MLX-dispatched graph for pt.logsumexp(x, axis=(1, 2)) on a 4-D input looks like:

Composite{(i0 + log(i1))}
 ├─ Max{axes=[1, 2]}(x)
 └─ Sum{axes=[1, 2]}
    └─ Composite{switch(isinf(max), exp(max), exp(x - max))}
       ├─ Isinf(ExpandDims(Max(x)))
       ├─ Exp(ExpandDims(Max(x)))
       ├─ x
       └─ ExpandDims(Max(x))

That's Max + ExpandDims + Exp + Isinf + Switch (Composite) + Sum + Composite{add(log)} — 7 separate dispatched calls plus the inf-guard work that MLX absorbs less efficiently than a single fused kernel.

MLX exposes mx.logsumexp(x, axis=...) as a native op. It is available as of MLX 0.31.2 (verified: getattr(mx, "logsumexp", None) is not None).

Reproducer + metrics

import time, statistics, mlx.core as mx, numpy as np
import pytensor, pytensor.tensor as pt

def sync(): mx.synchronize()
def bench(fn, *args, n=50, runs=5):
    sync()
    for _ in range(8):
        out = fn(*args); mx.eval(out)
    sync()
    out_t = []
    for _ in range(runs):
        sync(); t0 = time.perf_counter()
        for _ in range(n):
            out = fn(*args); mx.eval(out)
        sync()
        out_t.append((time.perf_counter() - t0) / n)
    return statistics.median(out_t) * 1e6

A, B, C, D = 4, 64, 64, 256
rng = np.random.default_rng(0)
x_mx = mx.array(rng.standard_normal((A, B, C, D)).astype(np.float32))
mx.eval(x_mx); sync()

@mx.compile
def raw_native(x): return mx.logsumexp(x, axis=(1, 2))

@mx.compile
def raw_manual(x):
    m = mx.max(x, axis=(1, 2), keepdims=True)
    return m.squeeze((1, 2)) + mx.log(mx.sum(mx.exp(x - m), axis=(1, 2)))

x_pt = pt.tensor4("x", dtype="float32")
f_pt = pytensor.function([x_pt], pt.logsumexp(x_pt, axis=(1, 2)), mode="MLX")
f_pt.trust_input = True

print(f"raw mx.logsumexp:        {bench(raw_native, x_mx):.1f} us")
print(f"raw manual (no guard):   {bench(raw_manual, x_mx):.1f} us")
print(f"pytensor MLX (current):  {bench(f_pt, x_mx):.1f} us")

→ 1.5× speedup by using mx.logsumexp. Saves ~33 % over the current PyTensor MLX path.

Proposed change

Two options, ordered by preference:

Option A (preferred): match the pattern at MLX rewrite time

Add an MLX-only rewrite that recognises the canonical logsumexp subgraph (Composite{add(log)} of Max and Sum(switch(isinf(max), exp(max), exp(x-max)))) and replaces it with a thin MLXLogSumExp(axis) Op that dispatches to mx.logsumexp. Keeps the inf-safety semantics because mx.logsumexp itself handles -inf correctly.

Files:

• pytensor/link/mlx/ops.py — MLXLogSumExp Op
• pytensor/link/mlx/dispatch/fused.py — dispatcher
• pytensor/link/mlx/rewriting/fused_logsumexp.py — pattern matcher

# dispatch
@mlx_funcify.register(MLXLogSumExp)
def mlx_funcify_MLXLogSumExp(op, **kwargs):
    axis = op.axis
    def logsumexp(x):
        return mx.logsumexp(x, axis=axis)
    return logsumexp

Option B (simpler, but heavier): teach pt.logsumexp to emit a recognisable Op

Add a LogSumExp(axis, keepdims) Op in pytensor/tensor/math.py that the helper builds when the user calls pt.logsumexp(...). Other backends keep the current Composite via perform/lowering rewrite; the MLX backend dispatches LogSumExp directly to mx.logsumexp.

This option is more invasive (touches pytensor/tensor/math.py and the public API), but avoids fragile pattern matching.

Acceptance criteria

• PyTensor pt.logsumexp graph compiled with mode="MLX" runs in ≤ 400 µs at the shape above (vs. 548 µs today).
• Numerics match mx.logsumexp within 1e-5, including for inputs containing -inf.
• Existing pt.logsumexp tests continue to pass on Numba/JAX backends (no regression for non-MLX modes).

Start by reading pytensor/link/mlx/ops.py, pytensor/link/mlx/dispatch/fused.py, and pytensor/link/mlx/rewriting/fused_logsumexp.py, then run the reproducer to inspect the current MLX graph and timing. Choose and implement one of the proposed approaches, verify the benchmark is at most 400 µs with matching -inf numerics, and run the existing logsumexp tests on Numba and JAX.