mx.sum over a non-last axis returns all zeros for certain large tensors (e.g. [1, 32, 4096, 131072])

Summary

mx.sum(x, axis=1) returns all zeros instead of the correct sum for x of shape [1, 32, 4096, 131072]. It happens deterministically, for every dtype (float32, float16, bfloat16). The same reduction over the last axis is correct, and smaller shapes are correct.

Repro

import mlx.core as mx

x = mx.ones((1, 32, 4096, 131072), dtype=mx.float32)
y = mx.sum(x, axis=1)
mx.eval(y)
print(float(y.min()), float(y.max()))   # prints 0.0 0.0, expected 32.0 32.0

100% reproducible (5 of 5 fresh runs). Same result with bfloat16 and float16.

What I found

Reducing the last axis is fine: mx.sum(x, axis=3) on the same tensor returns the correct 131072.0. Only the non last axis reduction is affected.

It is size gated. [1, 32, 2048, 131072] through [1, 32, 4032, 131072] are all correct. It first breaks at [1, 32, 4096, 131072] (2^34 elements). Just above that boundary the output is partially zeroed (for example [1, 32, 4352, 131072] gives min 0, max 32), which looks like a 32 bit index overflow in the strided reduction (2^34 is 4 times 2^32).

It is specific to a reduction code path, not a generic element count threshold. At the same 2^34 total, [1, 16, 8192, 131072], [1, 64, 2048, 131072], and [1, 128, 1024, 131072] are all correct. Only the shapes that reduce 32 elements with a large inner size break, e.g. [1, 32, 4096, 131072], [1, 32, 8192, 65536], [1, 32, 16384, 32768].

Expected vs actual

Expected: the correct per position sum (32.0 for the ones tensor above).
Actual: all zeros, or partially zeros just above the boundary.

Workaround

Reducing with elementwise adds (x[:, 0:1] + x[:, 1:2] + ...) or moving the reduced axis to the end both give the correct result.

Environment

MLX 0.32.0.dev20260609+cc3f3e60
Apple M3 Ultra, macOS 26.5.1 (build 25F80)

Start by running the reported Python mx.sum reproduction and compare non-last-axis reductions with the working last-axis case. Trace the reduction entry point for large strided tensors, then add a regression test covering [1, 32, 4096, 131072] and nearby shapes; done means the result is 32 for the ones tensor without partial zeros.