code: def sag_pool(
      gw, feature, ratio, graph_id, dataset, name, activation=L.tanh
  ):
      """Implementation of self-attention graph pooling (SAGPool)
  
      This is an implementation of the paper SELF-ATTENTION GRAPH POOLING
      (https://arxiv.org/pdf/1904.08082.pdf)
  
      Args:
          gw: Graph wrapper object.
  
          feature: A tensor with shape (num_nodes, feature_size).
  
          ratio: The pooling ratio of nodes we want to select.
  
          graph_id: The graphs that the nodes belong to.
  
          dataset: To differentiate FRANKENSTEIN dataset and other datasets.
  
          name: The name of SAGPool layer.
  
          activation: The activation function.
  
      Return:
          new_feature: A tensor with shape (num_nodes, feature_size), and the unselected
                       nodes' feature is masked by zero.
  
          ratio_length: The selected node numbers of each graph.
  
      """
      if dataset == "FRANKENSTEIN":
          gcn_ = gcn
      else:
          gcn_ = norm_gcn
  
      score = gcn_(
          gw=gw,
          feature=feature,
          hidden_size=1,
          activation=None,
          norm=gw.node_feat["norm"],
          name=name,
      )
      score = L.squeeze(score, axes=[])
      perm, ratio_length = topk_pool(gw, score, graph_id, ratio)
  
      mask = L.zeros_like(score)
      mask = L.cast(mask, dtype="float32")
      updates = L.ones_like(perm)
      updates = L.cast(updates, dtype="float32")
      mask = L.scatter(mask, perm, updates)
      new_feature = L.elementwise_mul(feature, mask, axis=0)
      temp_score = activation(score)
      new_feature = L.elementwise_mul(new_feature, temp_score, axis=0)
      return new_feature, 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 2 for calling a global function with example usage: sag_pool(gw, feature, ratio, graph_id, dataset, name, activation) and returns: new_feature,  ratio_length