[BUG] fft2/fftn over a length 1 axis can return wrong results on large Metal inputs #4704

☑️ I understand it is strictly prohibited to use AI to write issues.

Describe the bug

On Metal, fft2 over a length 1 axis sometimes gives wrong values (with no error) on the first call when the array is large, and then the right values after running it again. fft(x, axis=2) on the same array is always correct, and so is fft2 on a smaller array. This happened on 0.32.4 (main 36969c0) and on the 0.32.3 wheel.

To Reproduce

import mlx.core as mx

x = mx.ones((16, 1, 4194304), mx.complex64)
expected = mx.fft.fft(x, axis=2)
mx.eval(x, expected)
for i in range(4):
    out = mx.fft.fft2(x, axes=(1, 2))
    print(i, mx.array_equal(out, expected).item())

Expected behavior

0 True
1 True
2 True
3 True

Actual behavior

0 False
1 True
2 True
3 True

Each run was in a new process; I got these results 3/3 times on main (36969c0) and 2/3 times on the 0.32.3 wheel. (4, 1, 4194304) never failed.

Setting MLX_MAX_MB_PER_BUFFER=1 makes even mx.fft.fftn(mx.ones((1, 1048576), mx.complex64), axes=(0, 1)) fail on the first call, so I think the root cause relates to when the command buffer gets committed rather than the size directly.

Desktop

• OS Version: macOS 15.7.7
• Version: 0.32.4 (main 36969c0), also 0.32.3
• Apple M3 Pro

Additional context
I noticed this while working on #4702 , which is the same kind of bug. fft_op handles a length 1 axis by sharing the input buffer with the output (same dtype only), and nd_fft_op registers its intermediates as temporaries after the loop, so the final output can end up registered as a temporary and lose its fence. ifft2 processes the axes in the opposite order, so its length 1 step happens first and the input's buffer gets registered as a temporary.

I put up a draft PR (one open PR limit) with a fix in #4705.