code: def fft_shift(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
docs: 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.