ml-explore
Added Short time fourier transform and ISTFT implementations
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.
- [X] I have read the CONTRIBUTING document
- [X] I have run
pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes - [X] I have added tests that prove my fix is effective or that my feature works
- [ ] I have updated the necessary documentation (if needed)
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 would adding to the code to the utils.py be a better option ?
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).
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
@awni @angeloskath I have added a C++ implementation of stft in mlx.core.fft i.e. fft.cpp file. Can you please review it ??
The next step is to remove the Python one and make a binding in python/src/fft.cp
@awni I have added bindings and also removed the python implementation
@awni I fixed the headers. Compiled and tested it and it works fine on my end.