array-api
array-api copied to clipboard
Summary of FFT APIs
This issue gathers all the information of the current APIs in NumPy, CuPy, Dask, JAX, MXNet, pytorch and tensorflow in the FFT module. The main outcome of this issue is to open the discussion of the potential APIs that are going to be included in the spec.
The current APIs involved in this issue are:
- Standard FFT: fft, ifft, fft2, ifft2, fftn, ifftn
- Real FFT: rfft, irfft, rfft2, irfft2, rfftn, irfftn
- Hermitian FFT: hfft, ihfft
- Helper Routines: fftfreq, rfftfreq, fftshift, ifftshift
Individual APIs
fft/ifft
argument/keyword | NumPy | CuPy | Dask | JAX | MXNet | Pytorch | TF |
---|---|---|---|---|---|---|---|
n | Y | Y | Y | Y | N | Y | N |
axis/dim | Y | Y | Y | Y | N | Y | N |
norm | Y | Y | N | Y | N | Y | N |
data | N | N | N | N | Y | N | N |
out | N | N | N | N | Y | N | N |
name | N | N | N | N | Y | N | Y |
fft2/ifft2
argument/keyword | NumPy | CuPy | Dask | JAX | Pytorch | TF |
---|---|---|---|---|---|---|
s | Y | Y | Y | Y | Y | N |
axes/dim | Y | Y | Y | Y | Y | N |
norm | Y | Y | N | Y | Y | N |
name | N | N | N | N | N | Y |
These functions are not implemented in MXNet.
fftn/ifftn
argument/keyword | NumPy | CuPy | Dask | JAX | Pytorch |
---|---|---|---|---|---|
n | Y | Y | N | N | Y |
s | N | N | Y | Y | N |
axis/axes/dim | Y | Y | Y | Y | Y |
norm | Y | Y | N | Y | Y |
These functions are not implemented in MXNet and TF.
rfft/irfft
argument/keyword | NumPy | CuPy | Dask | JAX | Pytorch | TF |
---|---|---|---|---|---|---|
n | Y | Y | Y | Y | Y | N |
axis/dim | Y | Y | Y | Y | Y | N |
norm | Y | Y | N | Y | Y | N |
fft_length | N | N | N | N | N | Y |
name | N | N | N | N | N | Y |
These functions are not implemented in MXNet.
rfft2/irfft2
argument/keyword | NumPy | CuPy | Dask | JAX | Pytorch | TF |
---|---|---|---|---|---|---|
s | Y | Y | Y | Y | Y | N |
axes/dim | Y | Y | Y | Y | Y | N |
norm | Y | Y | N | Y | Y | N |
fft_length | N | N | N | N | N | Y |
name | N | N | N | N | N | Y |
These functions are not implemented in MXNet.
rfftn/irfftn
argument/keyword | NumPy | CuPy | Dask | JAX | Pytorch |
---|---|---|---|---|---|
s | Y | Y | Y | Y | Y |
axes/dim | Y | Y | Y | Y | Y |
norm | Y | Y | N | Y | Y |
These functions are not implemented in MXNet and TF.
hfft/ihfft
argument/keyword | NumPy | CuPy | Dask | JAX | Pytorch |
---|---|---|---|---|---|
n | Y | Y | Y | Y | Y |
axis/dim | Y | Y | Y | Y | Y |
norm | Y | Y | N | Y | Y |
These functions are not implemented in MXNet and TF.
fftfreq/rfftfreq
argument/keyword | NumPy | CuPy | Dask | JAX | Pytorch |
---|---|---|---|---|---|
d | Y | Y | Y | Y | Y |
dtype | N | N | N | N | Y |
device | N | N | N | N | Y |
requires_grad | N | N | N | N | Y |
layout | N | N | N | N | Y |
These functions are not implemented in MXNet and TF.
fftshift/ifftshift
argument/keyword | NumPy | CuPy | Dask | JAX | Pytorch |
---|---|---|---|---|---|
axes/dim | Y | Y | Y | Y | Y |
These functions are not implemented in MXNet and TF.
Summary
- Most FFT APIs have the following keyword arguments:
n
,axis
,norm
, with the exception of TF and MXNet. - The equivalent to the
axis
keyword in NumPy, CuPy, Dask and Jax isdim
in Pytorch. -
fft2
can have the same implementation offftn
if called with the correct arguments. -
fft3D
is only implemented in TF . - The output for the inverse functions return a complex array in NumPy.
- The arguments of the functions are the same for calculating the FFT and its inverse.
- There's a mix between the use of
axis
,axes
even among functions in the same library.
Thank you for the detailed research @steff456!
I haven't read this very carefully, just leaving a note. From the perspective of scientific computing, image and signal processing, etc., I'd like to make this appear asap (i.e., in the v1 draft) following the guideline for complex numbers (#153) so that participating libraries can start working on it if they already have support for complex numbers (like NumPy/CuPy/PyTorch), but maybe we need a discussion.
This looks great, thanks @steff456!
A couple of thoughts:
- We'd like to make this an "extension" rather than including it into the main API directly. Let's discuss how to document/specify those elsewhere, there's an issue open already somewhere.
- Let's keep in mind that there are several other implementations: PyFFTW,
scipy.fft
,mkl_fft
. CuPy also has it twice, one to match numpy and one to match scipy incupyx
- fortunately those APIs are identical.
I think norm
is just not currently implemented in Dask ( https://github.com/dask/dask/issues/3280 ), but I don't think it would be hard to add
To add to Ralf's point, I think there is an MPI FFT library. Not sure if we want to look at that as well
Also PyFFTW provides a Dask-based API. Though the API is the same as Dask's API
Based on the conversation in the meeting last week, we discussed:
-
scipy.fft
is missing from the table, but it's a superset of the NumPy API, with all extra functionalities like N-D Hermitian and R2R transforms (which probably have less interest among the library providers and users?) -
scipy.fft
has extra kw(?) arguments, for exampleoverwrite_x
. But its semantics is unclear and has been causing confusion over the years (ex: https://github.com/scipy/scipy/issues/10241, https://github.com/scipy/scipy/issues/8500), so if we're adopting it we must deviate from the current SciPy design (cc: @peterbell10, @grlee77) - @leofang proposed to not add
fft2
/ifft2
(and other 2D families), as it's just a special case forfftn
/ifftn
. No one objected, but it's best to solicit feedbacks from 3rd party libraries doing heavy imaging/signal processing.
- The output for the inverse functions return a complex array in NumPy.
It's easier to talk about the transform kinds, as "inverse" could be confusing. In short, the complex number support is a must for FFT, as they are used either as input (C2R), output (R2C), or both (C2C).
To add to Ralf's point, I think there is an MPI FFT library. Not sure if we want to look at that as well
AFAIK mpi4py-fft provides both a FFTW-like interface and a parallel interface on top of that. Most arguments are identical, with additional things due to the parallelism like MPI communicators and backends (fftw, numpy, scipy, ...). If there's strong interest we can move the discussion to the mpi4py-fft repo.
overwrite_x
I believe comes from the original f2py
wrapper code, or at least is consistent with that style. It was kept mainly for backwards-compatibility, I doubt it belongs in the array api standard.
Some corrections for the table:
-
fftn/ifftn
NumPy, CuPy and PyTorch are all documented as usings
notn
as the shape argument. - PyTorch does support
out
arguments for everything intorch.fft
exceptfftshift
andifftshift
. Although I notice it doesn't say that in the docs (my bad).
3. @leofang proposed to not add
fft2
/ifft2
(and other 2D families), as it's just a special case forfftn
/ifftn
. No one objected, but it's best to solicit feedbacks from 3rd party libraries doing heavy imaging/signal processing.
I checked in a few libraries (SciPy, cuCIM, scikit-image, scikit-learn, matplotlib) that 2D transforms (fft2
/ifft2
/rfft2
/irfft2
) are rarely used. At most one or two places in each codebase. I believe it is fine to not include them in the standard. Perhaps they were there because in the early days libraries wanted to be compatible with Matlab? (My guess)
Perhaps they were there because in the early days libraries wanted to be compatible with Matlab?
That's probably correct.
At most one or two places in each codebase. I believe it is fine to not include them in the standard.
Thanks for investigating @leofang. Your reasoning sounds good. And there's no issue with introduction or backwards compat here - the existing fft2
/ifft2
functions in libraries can stay as they are.
Hello'all, I just came up with a wild question while reviewing #189. I am curious if we can remove fft2
/ifft2
, is it really necessary then to special-case 1D transforms from N-D? I suppose it was done in the old days to provide a fast path for specialized codes, but shouldn't this purely be an implementation detail? I understand if we don't separate 1D from N-D, there'll likely be some boilerplate codes needed internally to distinguish them, but again it sounds like an implementation detail to me. Just a thought.
It is of course possible to express 1-D transforms in terms of n-D, but I suspect that if we make everyone write fftn(x, 1)
instead of fft(x)
, they won't be too happy. I think fft
is useful syntactic sugar, and it costs very little to support it.
This is an old issue - the fft
extension has been merged, which addressed this. So closing, thanks all.