[StockerMC]

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: I used Claude Code to help find this while writing Fix fast.cross_entropy VJP on non-contiguous Metal inputs #4702, verify the root cause, and write the fix and test. I reviewed and tested the changes myself.

fixes #4704

fft_op aliases the output onto the input for a size one axis. In nd_fft_op, this results in the primitive output sometimes sharing a buffer with an intermediate, and since the intermediates are then registered as temporaries, the output loses its fence the same way as #4702. With the inverse transform, the size one axis runs first so the input gets registered as a temporary instead.

Intermediates that share a buffer with the input or output are no longer registered as temporaries.

I added a regression test for both directions with an input large enough so that the command buffer splits after the FFT; the inverse case fails every time on 0.32.3 and 0.32.4, tested on an M3 Pro.

[zcbenz]

While this fixes the bug, we should fix the code that adds input/output to temp_arrs, rather than skipping them afterwards.

[StockerMC]

Thanks, I updated nd_fft_op to skip identity transforms before creating temporary buffers to prevent the aliases from happening at all.

[zcbenz]

Thanks for the update, just a note that I will need to land #4711 first which will add a deterministic check and can simplify the tests.