fft along a particular axis

[marcniethammer]

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?

[riceric22]

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.