Prune axis invariant offsets in (log)softmax

Description

import numpy as np, pytensor, pytensor.tensor as pt

N, G, K, M = 2000, 8, 10, 64
rng = np.random.default_rng(0)
x = pt.tensor("x", shape=(N, G, K))
c = pt.tensor("c", shape=(N, G, M))
offset = pt.exp(c).sum(-1, keepdims=True)  # (N, G, 1): constant along the softmax axis

# log_softmax(x + offset) == log_softmax(x): offset cancels, so computing it is dead work
slow = pytensor.function([x, c], pt.special.log_softmax(x + offset, -1))
fast = pytensor.function([x], pt.special.log_softmax(x, -1))

xv, cv = rng.normal(size=(N, G, K)), rng.normal(size=(N, G, M))
assert np.allclose(slow(xv, cv), fast(xv))

%timeit slow(xv, cv)   # 11.6 ms
%timeit fast(xv)       # 4.2 ms  -> 2.8x

Note our softmax and log_softmax already do the max subtraction stabilization, so a user doing the stabilization manually (which has the form of invariant addition/subtraction) can also be removed safely.

Similar logsumexp(x + c) -> logsumexp(x) + c

Start by reproducing the issue's Python example and confirming the invariant offset is removed without changing results. Then locate PyTensor's optimization handling for softmax, log_softmax, and logsumexp, and add coverage for axis-invariant additions or subtractions. Done means the optimized graph avoids the dead computation while preserving the shown numerical equivalences.