Add stft and istft to mlx.core.fft

[eyupcanakman]

Proposed changes

Adds stft and istft to mlx.core.fft, closes #1004. Both are written in C++ so they are available from every language binding.

stft frames the last axis with as_strided, applies the window, and runs rfft (or fft for two-sided output). istft inverts that with an overlap-add normalized by the squared-window envelope. The signature follows torch.stft, except norm takes the same "backward"/"ortho"/"forward" string the rest of mlx.fft uses. Batched input works and the output has shape (..., n_freq, n_frames).

Tested against torch.stft and torch.istft over reflect, constant, and edge padding, one-sided and two-sided transforms, the ortho norm, batched input, and n_fft that is not a multiple of hop_length. The C++ tests cover the shapes, an overlap-add round-trip, and the error paths.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[zcbenz]

Would #3608 help?

[eyupcanakman]

Yeah, its reflect mode matches what this helper does and also handles pad >= n, so I can drop the helper and call pad directly once it lands. Want me to rebase on it after it merges, or keep the local version for now?

[zcbenz]

Thanks for verifying, let's just keep the local version for now.

[zcbenz]

I must admit I'm not familiar with short-time fft at all, but the implementation looks good, and the tests are well-written, so this PR looks good to me.

[eyupcanakman]

Approved and green since mid-June. Not sure if this is waiting on something or just got buried.

[zcbenz]

@eyupcanakman Sorry for the long wait, I was waiting if someone else more familiar with fourier transform could take another look.

Would you mind replacing the pad_last_axis with pad op which supports reflect mode now? And then we can merge.

[eyupcanakman]

Replaced pad_last_axis with pad, removed the helper, and rebased on main (signals no longer than n_fft / 2 with reflect padding now follow NumPy's repeated reflection). Could you please approve the CI run when you get a chance?

[zcbenz]

Thanks for your patience!