Added Short time fourier transform and ISTFT implementations

[ParamThakkar123]

Proposed changes

Please include a description of the problem or feature this PR is addressing. If there is a corresponding issue, include the issue #.

This PR introduces the implementations of Short Time and Inverse Short Time Fourier Transforms in mlx as addressed in issue #1004

Fixes #1004

Checklist

Put an x in the boxes that apply.

• 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)

[ParamThakkar123]

I referred to librosa's documentation on stft to write my implementation:
https://librosa.org/doc/main/generated/librosa.stft.html

and Pytorch's documentation:
https://pytorch.org/docs/stable/generated/torch.stft.html

[angeloskath]

This looks well done from a first pass.

The main question to answer is whether this should be in the main mlx vs on a per library/use case implementation. I also think this shouldn't be in the top level of the package. We should come up with a better location if we are to merge this (I don't have a very good suggestion atm).

Regarding the PR content, I would definitely add docstrings and tests for the two functions.

[ParamThakkar123]

@angeloskath would adding to the code to the utils.py be a better option ?

[awni]

One option is to go the route of cloning scipy in MXL kind of like Jax (e.g. put this in a new package mlx.scipy.signal). It's a big package so I'm not really sure that's a good idea.

Another option is to put this in mlx.core.fft, but that would require a C++ implementation instead of a Python one (which may actually be handy to be honest because then we get it in all the APIs with simple bindings).

[ParamThakkar123]

One option is to go the route of cloning scipy in MXL kind of like Jax (e.g. put this in a new package mlx.scipy.signal). It's a big package so I'm not really sure that's a good idea.

Another option is to put this in mlx.core.fft, but that would require a C++ implementation instead of a Python one (which may actually be handy to be honest because then we get it in all the APIs with simple bindings).

Okay, I would implement a C++ version of this and add it to mlx.core.fft

[ParamThakkar123]

@awni should I create a separate namespace for this ??

[ParamThakkar123]

in fft.cpp

[ParamThakkar123]

@awni @angeloskath I have added a C++ implementation of stft in mlx.core.fft i.e. fft.cpp file. Can you please review it ??

[awni]

The next step is to remove the Python one and make a binding in python/src/fft.cp

[ParamThakkar123]

@awni I have added bindings and also removed the python implementation

[awni]

This PR is missing tests.

[ParamThakkar123]

@awni can you please review this PR ?

[ParamThakkar123]

I have added all the fixes and tests

[ParamThakkar123]

@awni I fixed the headers. Compiled and tested it and it works fine on my end.

[ParamThakkar123]

@awni @angeloskath

[ParamThakkar123]

@awni can you please review this PR ?

[ParamThakkar123]

@awni @angeloskath ???