class_docstr: Implementation of Graph Pooling "SAGPool"
      
  Reference Paper: Self-Attention Graph Pooling
  
  Args:
      input_dim: The size of the inputs.
      ratio: The ratio to reserve topk nodes.
      GNN: Graph Convolution Operation to generate score for every node. Default:None
      min_score: if min_score is not None, the operator only remove 
                          the node with value lower than min_score. Default:None
      nonlinearity: a nonlinear transform on score when min_score is not None. Default:None 
class_docstr_tok: ["Implementation","of","Graph","Pooling","``","SAGPool","''","Reference","Paper",":","Self-Attention","Graph","Pooling","Args",":","input_dim",":","The","size","of","the","inputs",".","ratio",":","The","ratio","to","reserve","topk","nodes",".","GNN",":","Graph","Convolution","Operation","to","generate","score","for","every","node",".","Default",":","None","min_score",":","if","min_score","is","not","None",",","the","operator","only","remove","the","node","with","value","lower","than","min_score",".","Default",":","None","nonlinearity",":","a","nonlinear","transform","on","score","when","min_score","is","not","None",".","Default",":","None"]
docstr: Args:          
      graph: `pgl.Graph` instance.
      x: A tensor with shape (num_nodes, input_size).
  Return:
      x: Reserved node features with shape (num_nodes * ratio, feature_size).
      batch: the updated graph_node_id.
      g:  updated 'pgl.Graph' instance (pgl.Graph).
docstr_tok: ["Args",":","graph",":","`","pgl.Graph","`","instance",".","x",":","A","tensor","with","shape","(","num_nodes",",","input_size",")",".","Return",":","x",":","Reserved","node","features","with","shape","(","num_nodes","*","ratio",",","feature_size",")",".","batch",":","the","updated","graph_node_id",".","g",":","updated","'pgl.Graph","'","instance","(","pgl.Graph",")","."]
code: def forward(self, graph, x):
      """
      Args:
          graph: `pgl.Graph` instance.
          x: A tensor with shape (num_nodes, input_size).
      Return:
          x: Reserved node features with shape (num_nodes * ratio, feature_size).
          batch: the updated graph_node_id.
          g:  updated 'pgl.Graph' instance (pgl.Graph).
      """
      batch = graph.graph_node_id
      attn_score = self.gnn(graph, x).reshape([-1])
      if self.min_score is None:
          attn_score = self.nonlinearity(attn_score)
      else:
          attn_score = segment_softmax(attn_score, batch)
      out, rank = segment_topk(
          x,
          attn_score,
          batch,
          self.ratio,
          self.min_score,
          return_index=True,
      )
      x = out * attn_score[rank].reshape([-1, 1])
      batch = batch[rank]
      edge_index, edge_attr = filter_adj(
          graph.edges, rank, num_nodes=attn_score.shape[0]
      )
      batch_size = graph.num_graph
      num_nodes = segment_sum(paddle.ones([x.shape[0]]), batch)
      graph_node_index = paddle.scatter(
          paddle.zeros([batch_size + 1]).astype(paddle.float32),
          paddle.arange(1, batch_size + 1),
          num_nodes.cumsum(0).astype(paddle.float32),
      ).astype(paddle.int64)
      g = pgl.Graph(
          num_nodes=x.shape[0],
          edges=edge_index,
          node_feat={"attr": x},
          _graph_node_index=graph_node_index,
          _num_graph=(batch.max() + 1).item(),
      )
  
      # g = g.tensor()
      return x, batch, g
  
imports: import warnings
  import numpy
  import paddle
  import paddle.nn
  import paddle.nn.functional
  import pgl
  from pgl.math import segment_sum, segment_softmax, segment_topk, segment_pool
  from pgl.utils.transform import filter_adj
label_desc: Use image node_id 2 for calling the SAGPool obj's underlying member method code with example usage: obj.forward(graph, x) and returns: x,  batch,  g