code: def gather(x, idx, method=2):
      """
      https://github.com/pytorch/pytorch/issues/15245
      implementation of a custom gather operation for faster backwards.
      :param x: input with shape [N, D_1, ... D_d]
      :param idx: indexing with shape [n_1, ..., n_m]
      :param method: Choice of the method
      :return: x[idx] with shape [n_1, ..., n_m, D_1, ... D_d]
      """
      idx[idx == -1] = x.shape[0] - 1  # Shadow point
      if method == 0:
          return x[idx]
      elif method == 1:
          x = x.unsqueeze(1)
          x = x.expand((-1, idx.shape[-1], -1))
          idx = idx.unsqueeze(2)
          idx = idx.expand((-1, -1, x.shape[-1]))
          return x.gather(0, idx)
      elif method == 2:
          for i, ni in enumerate(idx.size()[1:]):
              x = x.unsqueeze(i + 1)
              new_s = list(x.size())
              new_s[i + 1] = ni
              x = x.expand(new_s)
          n = len(idx.size())
          for i, di in enumerate(x.size()[n:]):
              idx = idx.unsqueeze(i + n)
              new_s = list(idx.size())
              new_s[i + n] = di
              idx = idx.expand(new_s)
          return x.gather(0, idx)
      else:
          raise ValueError("Unkown method")
  
label_desc: Use image node_id 1 for calling a global function with example usage: gather(x, idx, method) and returns: x, x, x