docs: def convolve_spatial2(im, psfs,
                        mode = "constant",
                        grid_dim = None,
                        sub_blocks = None,
                        pad_factor = 2,
                        plan = None,
                        return_plan = False):
      """
      GPU accelerated spatial varying convolution of an 2d image with a
      (Gy,Gx) grid of psfs assumed to be equally spaced within the image
      the input image im is subdivided into (Gy,Gx) blocks, each block is
      convolved with the corresponding psf and linearly interpolated to give the
      final result
  
      The psfs can be given either in
  
      A) Stackmode
      psfs.shape =  (Gy,Gx, Hy, Hx)
      then psfs[j,i] is the psf at the center of each block (i,j) in the image
  
      B) Flatmode
      psfs.shape = im.shape
      then the psfs are assumed to be definied on the gridpoints of the images itself
      in this case grid_dim = (Gy,Gx) has to be given
  
      as of now each image dimension has to be divisible by the grid dim, i.e.
      ::
          Nx % Gx == 0
          Ny % Gy == 0
  
      GPU Memory consumption is of order 8*Nx*Ny
      If not enough GPU memory is available, consider using
      sub_blocks = (n,m)
      then the operation is carried out in a tiled fashion reducing
      memory consumption to 8*Nx*Ny*(1/n+2/Gx)*(1/m+2/Gy)
  
      Example
      -------
  
  
      im = np.zeros((128,128))
      im[::10,::10] = 1.
  
      # Stackmode
      psfs = np.ones((16,16,4,4))
      res = convolve_spatial2(im, psfs, mode = "wrap")
  
      # Flatmode
      _X,_Y = np.meshgrid(*(np.arange(128),)*2)
      psfs = np.clip(np.sin(2*np.pi*_X/8),0,1)*np.clip(np.cos(2*np.pi*_Y/8),0,1)
      res = convolve_spatial2(im, psfs, grid_dim = (16,16))
  
  
      Parameters
      ----------
      im: ndarray
          the image to convolve
      psfs: ndarray
          the (Gx,Gy) psf grid, either of shape (Gx,Gy, Hy, Hx) or im.shape
  
      mode: string, optional
          padding mode, either "constant" or "wrap"
      grid_dim: tuple, optional
          the (Gy,Gx) grid dimension, has to be provided if psfs.shape = im.shape
  
      sub_blocks: tuple, optional
          tiling mode, give e.g. (2,2) to sequentially operate on quadratnts
      pad_factor: int
          the factor of its size each block get tiled, use pad_factor=2 if the psfs
          are well localized, use pad_factor = 3 if not (e.g. if you experience blocking)_
  
      plan: fft_plan, optional
          when given use this as the fft plan
      return_plan: bool, optional
          return (res, plan) with plan being the fft plan for further use
  
      Returns
      -------
      res: ndarray
          the convolved image
  
  
  
      """
  
      ndim = im.ndim
  
      if ndim != 2:
          raise ValueError("wrong dimensions of input!")
  
  
      if grid_dim:
          if psfs.shape != im.shape:
              raise ValueError("if grid_dim is set, then im.shape = hs.shape !")
      else:
          if not psfs.ndim==2*ndim:
              raise ValueError("wrong dimensions of psf grid! (Gy,Gx,Ny,Nx)")
  
      if grid_dim:
          Gs = grid_dim
      else:
          Gs = psfs.shape[:ndim]
  
      if not np.all([n%g==0 for n,g in zip(im.shape,Gs)]):
          raise NotImplementedError("shape of image has to be divisible by Gx Gy  = %s shape mismatch"%(str(psfs.shape[:2])))
  
      if sub_blocks is None:
          return _convolve_spatial2(im,
                                    psfs,
                                    mode = mode,
                                    pad_factor = pad_factor,
                                    plan = plan,
                                    return_plan = return_plan,
                                    grid_dim = grid_dim)
      else:
          # cut the image into tile and operate on every of them
          N_sub = [n//s for n,s  in zip(im.shape,sub_blocks)]
          Nblocks = [n//g for n,g  in zip(im.shape,Gs)]
          Npads = [n*(s>1) for n,s  in zip(Nblocks, sub_blocks)]
          grid_dim_sub = [g//s+2*(s>1) for g,s   in zip(Gs, sub_blocks)]
  
          logger.debug(
              "N_sub: {}, Nblocks: {}, grid_dim_sub, {}, Npads, {}".format(
                  N_sub, Nblocks, grid_dim_sub, Npads
              ))
  
          if grid_dim:
              res = np.empty(im.shape, np.float32)
              plan = None
              for (im_tile, im_s_src, im_s_dest), (hs_tile, hs_s_src, hs_s_dest)\
                  in zip(tile_iterator(im,blocksize=N_sub,
                                       padsize=Npads,
                                       mode = mode),\
                      tile_iterator(psfs, blocksize=N_sub,
                                    padsize=Npads,
                                    mode = mode)):
  
                  res_tile, plan = _convolve_spatial2(im_tile.copy(),
                                                hs_tile.copy(),
                                                mode = mode,
                                                pad_factor = pad_factor,
                                                return_plan=True,
                                                plan = plan,
                                                grid_dim = grid_dim_sub)
  
  
                  res[im_s_src] = res_tile[im_s_dest]
              return res
  
          else:
              raise NotImplementedError()