code: def topk(
      x: torch.Tensor,
      ratio: float,
      batch_id: torch.Tensor,
      num_nodes: torch.Tensor,
  ):
      """The top-k pooling method. Given a graph batch, this method will pool out some
      nodes from input node feature tensor for each graph according to the given ratio.
  
      Args:
          x (torch.Tensor): The input node feature batch-tensor to be pooled.
          ratio (float): the pool ratio. For example if :obj:`ratio=0.5` then half of the input
              tensor will be pooled out.
          batch_id (torch.Tensor): The batch_id of each element in the input tensor.
          num_nodes (torch.Tensor): The number of nodes of each graph in batch.
  
      Returns:
          perm (torch.Tensor): The index in batch to be kept.
          k (torch.Tensor): The remaining number of nodes for each graph.
      """
      batch_size, max_num_nodes = (
          num_nodes.size(0),
          num_nodes.max().item(),
      )
  
      cum_num_nodes = torch.cat(
          [num_nodes.new_zeros(1), num_nodes.cumsum(dim=0)[:-1]], dim=0
      )
  
      index = torch.arange(
          batch_id.size(0), dtype=torch.long, device=x.device
      )
      index = (index - cum_num_nodes[batch_id]) + (
          batch_id * max_num_nodes
      )
  
      dense_x = x.new_full(
          (batch_size * max_num_nodes,), torch.finfo(x.dtype).min
      )
      dense_x[index] = x
      dense_x = dense_x.view(batch_size, max_num_nodes)
  
      _, perm = dense_x.sort(dim=-1, descending=True)
      perm = perm + cum_num_nodes.view(-1, 1)
      perm = perm.view(-1)
  
      k = (ratio * num_nodes.to(torch.float)).ceil().to(torch.long)
      mask = [
          torch.arange(k[i], dtype=torch.long, device=x.device)
          + i * max_num_nodes
          for i in range(batch_size)
      ]
  
      mask = torch.cat(mask, dim=0)
      perm = perm[mask]
  
      return perm, k
  
imports: import logging
  import math
  import torch
  from scipy.stats import t
label_desc: Use image node_id 3 for calling a global function with example usage: topk(x, ratio, batch_id, num_nodes) and returns: perm,  k