Add fftfreq, rfftfreq and scalar axes for fftshift/ifftshift

[declanhealy2]

Proposed changes

Add fftfreq and rfftfreq to mlx.core.fft, and allow fftshift/ifftshift to accept a scalar int for the axes parameter to match NumPy's API.

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)

[angeloskath]

That looks good but can you add it to the C++ side as well and then call it from python?

[declanhealy2]

Thanks for looking at this. Have now moved fftfreq and rfftfreq to the C++ side.

[angeloskath]

Looks great. I 'll merge after the tests pass.