code: def topk_pool(gw, score, graph_id, ratio):
      """Implementation of topk pooling, where k means pooling ratio.
  
      Args:
          gw: Graph wrapper object.
  
          score: The attention score of all nodes, which is used to select
                 important nodes.
  
          graph_id: The graphs that the nodes belong to.
  
          ratio: The pooling ratio of nodes we want to select.
  
      Return:
          perm: The index of nodes we choose.
  
          ratio_length: The selected node numbers of each graph.
      """
  
      graph_lod = gw.graph_lod
      graph_nodes = gw.num_nodes
      num_graph = gw.num_graph
  
      num_nodes = L.ones(shape=[graph_nodes], dtype="float32")
      num_nodes = L.lod_reset(num_nodes, graph_lod)
      num_nodes_per_graph = L.sequence_pool(num_nodes, pool_type="sum")
      max_num_nodes = L.reduce_max(num_nodes_per_graph, dim=0)
      max_num_nodes = L.cast(max_num_nodes, dtype="int32")
  
      index = L.arange(0, gw.num_nodes, dtype="int64")
      offset = L.gather(graph_lod, graph_id, overwrite=False)
      index = (index - offset) + (graph_id * max_num_nodes)
      index.stop_gradient = True
  
      # padding
      dense_score = L.fill_constant(
          shape=[num_graph * max_num_nodes],
          dtype="float32",
          value=-999999,
      )
      index = L.reshape(index, shape=[-1])
      dense_score = L.scatter(dense_score, index, updates=score)
      num_graph = L.cast(num_graph, dtype="int32")
      dense_score = L.reshape(
          dense_score, shape=[num_graph, max_num_nodes]
      )
  
      # record the sorted index
      _, sort_index = L.argsort(dense_score, axis=-1, descending=True)
  
      # recover the index range
      graph_lod = graph_lod[:-1]
      graph_lod = L.reshape(graph_lod, shape=[-1, 1])
      graph_lod = L.cast(graph_lod, dtype="int64")
      sort_index = L.elementwise_add(sort_index, graph_lod, axis=-1)
      sort_index = L.reshape(sort_index, shape=[-1, 1])
  
      # use sequence_slice to choose selected node index
      pad_lod = L.arange(
          0,
          (num_graph + 1) * max_num_nodes,
          step=max_num_nodes,
          dtype="int32",
      )
      sort_index = L.lod_reset(sort_index, pad_lod)
      ratio_length = L.ceil(num_nodes_per_graph * ratio)
      ratio_length = L.cast(ratio_length, dtype="int64")
      ratio_length = L.reshape(ratio_length, shape=[-1, 1])
      offset = L.zeros(shape=[num_graph, 1], dtype="int64")
      choose_index = L.sequence_slice(
          input=sort_index, offset=offset, length=ratio_length
      )
  
      perm = L.reshape(choose_index, shape=[-1])
      return perm, ratio_length
  
file_path: PGL/legacy/examples/SAGPool/layers.py
imports: import numpy
  import paddle
  import paddle.fluid
  import paddle.fluid.layers
  import pgl
  from pgl.graph_wrapper import GraphWrapper
  from pgl.utils.logger import log
  from conv import norm_gcn
  from pgl.layers.conv import gcn
label_desc: Use image node_id 1 for calling a global function with example usage: topk_pool(gw, score, graph_id, ratio) and returns: perm,  ratio_length