pytorch_fft
pytorch_fft copied to clipboard
fft along a particular axis
In numpy's fft functions it is possible to perform ffts with respect to a particular axis: https://docs.scipy.org/doc/numpy-1.13.0/reference/routines.fft.html
Is it possibly to add this to the code? Or is it already in there somehow?
CuFFT requires that the data follow their 'advanced data layout' specification, which to my knowledge, limits the ability to perform FFT to an arbitrary axis. The easiest way I can think of is to use the PyTorch permute function to move the desired FFT axis, which matches existing CuFFT approaches.