2D FFTs via cuFFTDx are implemented via several 1D FFTs as shown here: https://github.com/NVIDIA/CUDALibrarySamples/blob/master/MathDx/cuFFTDx/05_fft_Xd/fft_2d.cu. This issue is to augment the current JIT support in the FFT operator to support 2D where it can.