Add this to speedup and reduce memory use when computing logsumexp during e.g. lora fine tuning. The large vocab (200k+ in some cases) makes this optimization important. Also since the logsumexp is in high precision this let's us avoid up casting the logits and instead only upcasting the per-token loss prior to the reduction.

QLoRA fine tuning Gemma 3 1B:

Pre: Iter 10: Train loss 3.595, Learning Rate 1.000e-05, It/sec 0.227, Tokens/sec 812.842, Trained Tokens 35845, Peak mem 66.335 GB

Post: Iter 10: Train loss 3.594, Learning Rate 1.000e-05, It/sec 0.345, Tokens/sec 1235.568, Trained Tokens 35845, Peak mem 49.164 GB

And a microbenchmark:

shape = (4096, 120_000)

x = mx.random.uniform(shape=shape)

def fun(x):
    return [mx.logsumexp(x, axis=-1, keepdims=True) for _ in range(100)]

Pre: 1588.345
Post: 257.241