https://github.com/alasdairtran/fourierflow/blob/97e6cfb0848e44d3a7bc1d063b1ab86bc4c603ee/fourierflow/modules/fourier_2d_factorized_parallel.py#L218