code: def graph_gather(gw, feature, index):
      """Implementation of graph gather
  
      Gather the corresponding index for each graph.
  
      Args:
          gw: Graph wrapper object (:code:`StaticGraphWrapper` or :code:`GraphWrapper`)
  
          feature: A tensor with shape (num_nodes, ).
  
          index (int32): A tensor with K-rank where the first dim denotes the graph.
                          Shape (num_graph, ) or (num_graph, k1, k2, k3, ..., kn).
                         WARNING: We dont support negative index.
  
      Return:
          A tensor with shape (num_graph, k1, k2, k3, ..., kn, hidden_size)
      """
      shape = L.shape(index)
      output_dim = int(feature.shape[-1])
      index = index + gw.graph_lod[:-1]
      index = L.reshape(index, [-1])
      feature = L.gather(feature, index, overwrite=False)
      new_shape = []
      for i in range(shape.shape[0]):
          new_shape.append(shape[i])
      new_shape.append(output_dim)
      feature = L.reshape(feature, new_shape)
      return feature
  
file_path: PGL/legacy/pgl/layers/graph_op.py
imports: import paddle.fluid
  import paddle.fluid.layers
  from pgl import graph_wrapper
  from pgl.utils import paddle_helper
  from pgl.utils import op
label_desc: Use image node_id 3 for calling a global function with example usage: graph_gather(gw, feature, index) and returns: feature