This function performs the FFT shift operation on the GPU to restore the correct output shape.

    INPUTS:
    in1     (no default):   Array containing data which has been FFTed and IFFTed.

    OUTPUTS:
    in1                     FFT-shifted version of in1.

def fft_shift(in1):
    """
    This function performs the FFT shift operation on the GPU to restore the correct output shape.

    INPUTS:
    in1     (no default):   Array containing data which has been FFTed and IFFTed.

    OUTPUTS:
    in1                     FFT-shifted version of in1.
    """

    ker = SourceModule("""
                        __global__ void fft_shift_ker(float *in1)
                        {
                            const int len = gridDim.x*blockDim.x;
                            const int col = (blockDim.x * blockIdx.x + threadIdx.x);
                            const int row = (blockDim.y * blockIdx.y + threadIdx.y);
                            const int tid2 = col + len*row;
                            float tmp = 0;

                            if ((col<(len/2))&(row<(len/2)))
                                {
                                    tmp = in1[tid2];
                                    in1[tid2] = in1[(len/2)*(len + 1) + tid2];
                                    in1[(len/2)*(len + 1) + tid2] = tmp;
                                }

                            if ((col>=(len/2))&(row<(len/2)))
                                {
                                    tmp = in1[tid2];
                                    in1[tid2] = in1[(len/2)*(len - 1) + tid2];
                                    in1[(len/2)*(len - 1) + tid2] = tmp;
                                }

                        }
                       """, keep=True)

    fft_shift_ker = ker.get_function("fft_shift_ker")
    fft_shift_ker(in1, block=(32,32,1), grid=(int(in1.shape[1]//32), int(in1.shape[0]//32)))

    return in1