Accumulate low-precision float reductions in float32 on the CPU

[v-code01]

Problem

Sum / Prod reductions of float16 / bfloat16 accumulate in the input dtype (reduction_op<InT, InT, ...>), so a long reduction saturates the accumulator:

with mx.stream(mx.cpu):
    mx.sum(mx.ones((20000,), mx.float16))    # 16384.0  (true 20000, numpy: 20000)
    mx.sum(mx.ones((70000,), mx.bfloat16))   #   256.0  (true 70000)
    mx.mean(mx.ones((70000,), mx.bfloat16))  #     0.0  (true 1.0)

The CPU result is both wrong and worse than numpy, which reduces in a wider type. Because mean and var are built on sum, they inherit the error (a large bfloat16 mean of ones returns 0). This is the reduce-path sibling of the scan issue in #3907.

Fix

Integer reductions already widen to int32; do the same for low-precision floats. Reduce into a float32 accumulator through the existing reduction_op machinery — which already supports an accumulator wider than the input (used for int8/int16 -> int32) — and narrow back to the output dtype once at the end. Every other dtype is untouched, so float32 / int / complex results stay bit-identical.

with mx.stream(mx.cpu):
    mx.sum(mx.ones((20000,), mx.float16))    # 20000.0
    mx.sum(mx.ones((70000,), mx.bfloat16))   # 70144.0  (nearest bfloat16)
    mx.mean(mx.ones((70000,), mx.bfloat16))  #     1.0

A float16 sum still overflows to inf once the true value exceeds float16's ~65504 max — that is unavoidable for a float16 output — but within range it now matches numpy, and mean/var are correct.

Tests

Added test_low_precision_float_accumulation (float16/bfloat16 sum/mean/var, all-reduce and a strided axis reduction), pinned to the CPU stream. It fails on the current backend (stall/saturation) and passes with the fix. Full test_reduce.py + test_ops.py (157 tests / 5371 subtests) stay green; float32/int reductions remain bit-identical to numpy.

Verified with a CPU-only build (MLX_BUILD_METAL=OFF) on an M4.

[jagrit06]

Thanks for putting this together
Have you had a chance to run the CPU reduction benchmarks to see if the change has any implications on performance ?