-
Notifications
You must be signed in to change notification settings - Fork 118
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add FFT Ops #480
Add FFT Ops #480
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This looks good to me, just a few comments
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the PR!
Are you able to add it to the numpy backend as well? We can probably just redirect to jax like we do for most nn ops. |
Redirected to JAX for NumPy backend! |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM -- thanks for the great contribution! 👍
* Add FFT Ops * Fixes * Fix torch * Address Matt's comments * Address Francois' comments * Shift docstrings to correct fns * Add NumPy backend FFT ops * Fix numpy backend * Minor change * Redirect NumPy FFT to JAX
Resolves #429
The reason for not supporting all args at the moment is that TensorFlow supports none of those args. Moreover, there are some inconsistencies between the NumPy FFT API and JAX FFT API. We need a very minimal version for FNet (and for most usecases, I suppose). This can be extended later to support all args.