code: import torch
  from torch import nn
  from torch.nn import CrossEntropyLoss
  import logging
  from .bert_model import BertPreTrainedModel, BertPreTrainingHeads, BertModel, BertEncoder, BertPooler, BertLayerNorm
  
  
  logger = logging.getLogger(__name__)
  
  
  class DualPositionBertEmbeddings(nn.Module):
      """Construct the embeddings from word, position and token_type embeddings."""
  
      def __init__(self, config):
          super(DualPositionBertEmbeddings, self).__init__()
          self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
          self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
          self.position_embeddings_second = nn.Embedding(config.max_position_embeddings, config.hidden_size)
          self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
          # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
          # any TensorFlow checkpoint file
          self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-12)
          self.dropout = nn.Dropout(config.hidden_dropout_prob)
  
      def forward(self, input_ids, token_type_ids, position_ids, position_ids_second):
          if token_type_ids is None:
              token_type_ids = torch.zeros_like(input_ids)
  
          words_embeddings = self.word_embeddings(input_ids)
          position_embeddings = self.position_embeddings(position_ids)
          position_embeddings_second = self.position_embeddings(position_ids_second)
          token_type_embeddings = self.token_type_embeddings(token_type_ids)
  
          embeddings = words_embeddings + position_embeddings + position_embeddings_second + token_type_embeddings
          embeddings = self.LayerNorm(embeddings)
          embeddings = self.dropout(embeddings)
          return embeddings
  
  
  class DualPositionBertModel(BertModel):
      def __init__(self, config):
          super(DualPositionBertModel, self).__init__(config)
          self.embeddings = DualPositionBertEmbeddings(config)
          self.encoder = BertEncoder(config)
          self.pooler = BertPooler(config)
          self.apply(self.init_bert_weights)
          logger.info("Init BERT pretrain model")
  
      def forward(
          self,
          input_ids,
          token_type_ids=None,
          attention_mask=None,
          output_all_encoded_layers=True,
          checkpoint_activations=False,
          position_ids=None,
          position_ids_second=None,
      ):
          if attention_mask is None:
              attention_mask = torch.ones_like(input_ids)
          if token_type_ids is None:
              token_type_ids = torch.zeros_like(input_ids)
  
          if len(attention_mask.shape) == 2:
              extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
          elif len(attention_mask.shape) == 3:
              extended_attention_mask = attention_mask.unsqueeze(1)
          else:
              raise Exception("invalid attention mask shape! shape: %s" % (attention_mask.shape))
  
          extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
          extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
          embedding_output = self.embeddings(input_ids, token_type_ids, position_ids, position_ids_second)
          encoded_layers = self.encoder(
              embedding_output,
              extended_attention_mask,
              output_all_encoded_layers=output_all_encoded_layers,
              checkpoint_activations=checkpoint_activations,
          )
          sequence_output = encoded_layers[-1]
          pooled_output = self.pooler(sequence_output)
  
          if not output_all_encoded_layers:
              encoded_layers = encoded_layers[-1]
          return encoded_layers, pooled_output
  
  
  class DualPositionBertForPreTrainingPreLN(BertPreTrainedModel):
      """BERT model with pre-training heads and dual position
      Params:
          config: a BertConfig class instance with the configuration to build a new model.
  
      """
  
      def __init__(self, config):
          super(DualPositionBertForPreTrainingPreLN, self).__init__(config)
          self.bert = DualPositionBertModel(config)
          self.cls = BertPreTrainingHeads(config, self.bert.embeddings.word_embeddings.weight)
          self.apply(self.init_bert_weights)
  
      def forward(
          self,
          input_ids,
          token_type_ids=None,
          attention_mask=None,
          masked_lm_labels=None,
          position_ids=None,
          position_ids_second=None,
          log=True,
      ):
          sequence_output, pooled_output = self.bert(
              input_ids=input_ids,
              token_type_ids=token_type_ids,
              attention_mask=attention_mask,
              output_all_encoded_layers=False,
              checkpoint_activations=False,
              position_ids=position_ids,
              position_ids_second=position_ids_second,
          )
  
          if masked_lm_labels is not None:
              # filter out all masked labels.
              masked_token_indexes = torch.nonzero((masked_lm_labels + 1).view(-1)).view(-1)
              prediction_scores, _ = self.cls(sequence_output, pooled_output, masked_token_indexes)
              target = torch.index_select(masked_lm_labels.view(-1), 0, masked_token_indexes)
  
              loss_fct = CrossEntropyLoss(ignore_index=-1)
              masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), target)
              return masked_lm_loss
          else:
              prediction_scores, _ = self.cls(sequence_output, pooled_output)
              return prediction_scores
   from typing import Union, Callable
  import numpy as np
  import warnings
  
  import torch
  import torch.nn as nn
  
  from sklearn.metrics import f1_score
  
  
  def setup_evaluator(metric: Union[str, Callable]):
      if isinstance(metric, str):
          metric = metric.lower()
          if metric == "acc" or metric == "accuracy":
              return Accuracy()
          elif metric == "multilabel_microf1" or "microf1" or "micro_f1":
              return MultiLabelMicroF1()
          elif metric == "multiclass_microf1":
              return MultiClassMicroF1()
          else:
              raise NotImplementedError
      else:
          return BaseEvaluator(metric)
  
  
  class BaseEvaluator(object):
      def __init__(self, eval_func):
          self.y_pred = list()
          self.y_true = list()
          self.eval_func = eval_func
  
      def __call__(self, y_pred, y_true):
          metric = self.eval_func(y_pred, y_true)
          self.y_pred.append(y_pred.cpu())
          self.y_true.append(y_true.cpu())
          return metric
  
      def clear(self):
          self.y_pred = list()
          self.y_true = list()
  
      def evaluate(self):
          if len(self.y_pred) > 0:
              y_pred = torch.cat(self.y_pred, dim=0)
              y_true = torch.cat(self.y_true, dim=0)
              self.clear()
              return self.eval_func(y_pred, y_true)
          return 0
  
  
  class Accuracy(object):
      def __init__(self, mini_batch=False):
          super(Accuracy, self).__init__()
          self.mini_batch = mini_batch
          self.tp = list()
          self.total = list()
  
      def __call__(self, y_pred, y_true):
          pred = (y_pred.argmax(1) == y_true).int()
          tp = pred.sum().int()
          total = pred.shape[0]
          if torch.is_tensor(tp):
              tp = tp.item()
  
          # if self.mini_batch:
          self.tp.append(tp)
          self.total.append(total)
  
          return tp / total
  
      def evaluate(self):
          if len(self.tp) > 0:
              tp = np.sum(self.tp)
              total = np.sum(self.total)
              self.tp = list()
              self.total = list()
              return tp / total
          warnings.warn("pre-computing list is empty")
          return 0
  
      def clear(self):
          self.tp = list()
          self.total = list()
  
  
  class MultiLabelMicroF1(Accuracy):
      def __init__(self, mini_batch=False):
          super(MultiLabelMicroF1, self).__init__(mini_batch)
  
      def __call__(self, y_pred, y_true, sigmoid=False):
          if sigmoid:
              border = 0.5
          else:
              border = 0
          y_pred[y_pred >= border] = 1
          y_pred[y_pred < border] = 0
          tp = (y_pred * y_true).sum().to(torch.float32).item()
          fp = ((1 - y_true) * y_pred).sum().to(torch.float32).item()
          fn = (y_true * (1 - y_pred)).sum().to(torch.float32).item()
          total = tp + fp + fn
  
          # if self.mini_batch:
          self.tp.append(int(tp))
          self.total.append(int(total))
  
          if total == 0:
              return 0
          return float(tp / total)
  
  
  class MultiClassMicroF1(Accuracy):
      def __init__(self, mini_batch=False):
          super(MultiClassMicroF1, self).__init__(mini_batch)
  
  
  class CrossEntropyLoss(nn.Module):
      def __call__(self, y_pred, y_true):
          y_true = y_true.long()
          y_pred = torch.nn.functional.log_softmax(y_pred, dim=-1)
          return torch.nn.functional.nll_loss(y_pred, y_true)
  
  
  class BCEWithLogitsLoss(nn.Module):
      def __call__(self, y_pred, y_true, reduction="mean"):
          y_true = y_true.float()
          loss = torch.nn.BCEWithLogitsLoss(reduction=reduction)(y_pred, y_true)
          if reduction == "none":
              loss = torch.sum(torch.mean(loss, dim=0))
          return loss
  
  
  def multilabel_f1(y_pred, y_true, sigmoid=False):
      if sigmoid:
          y_pred[y_pred > 0.5] = 1
          y_pred[y_pred <= 0.5] = 0
      else:
          y_pred[y_pred > 0] = 1
          y_pred[y_pred <= 0] = 0
      tp = (y_true * y_pred).sum().to(torch.float32)
      # tn = ((1 - y_true) * (1 - y_pred)).sum().to(torch.float32)
      fp = ((1 - y_true) * y_pred).sum().to(torch.float32)
      fn = (y_true * (1 - y_pred)).sum().to(torch.float32)
  
      epsilon = 1e-7
      precision = tp / (tp + fp + epsilon)
      recall = tp / (tp + fn + epsilon)
      f1 = (2 * precision * recall) / (precision + recall + epsilon)
      return f1.item()
  
  
  def multiclass_f1(y_pred, y_true):
      y_true = y_true.squeeze().long()
      preds = y_pred.max(1)[1]
      preds = preds.cpu().detach().numpy()
      labels = y_true.cpu().detach().numpy()
      micro = f1_score(labels, preds, average="micro")
      return micro
  
  
  def accuracy(y_pred, y_true):
      y_true = y_true.squeeze().long()
      preds = y_pred.max(1)[1].type_as(y_true)
      correct = preds.eq(y_true).double()
      correct = correct.sum().item()
      return correct / len(y_true)
  
  
  def cross_entropy_loss(y_pred, y_true):
      y_true = y_true.long()
      y_pred = torch.nn.functional.log_softmax(y_pred, dim=-1)
      return torch.nn.functional.nll_loss(y_pred, y_true)
  
  
  def bce_with_logits_loss(y_pred, y_true, reduction="mean"):
      y_true = y_true.float()
      loss = torch.nn.BCEWithLogitsLoss(reduction=reduction)(y_pred, y_true)
      if reduction == "none":
          loss = torch.sum(torch.mean(loss, dim=0))
      return loss
   import numpy as np
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from scipy.linalg import block_diag
  
  from cogdl.layers import SAGELayer
  from cogdl.utils import split_dataset_general
  
  from .. import BaseModel
  
  
  class EntropyLoss(nn.Module):
      # Return Scalar
      def forward(self, adj, anext, s_l):
          # entropy.mean(-1).mean(-1): 1/n in node and batch
          # entropy = (torch.distributions.Categorical(
          # probs=s_l).entropy()).sum(-1).mean(-1)
          entropy = (torch.distributions.Categorical(probs=s_l).entropy()).mean()
          assert not torch.isnan(entropy)
          return entropy
  
  
  class LinkPredLoss(nn.Module):
      def forward(self, adj, anext, s_l):
          link_pred_loss = (adj - s_l.matmul(s_l.transpose(-1, -2))).norm(dim=(1, 2))
          link_pred_loss = link_pred_loss / (adj.size(1) * adj.size(2))
          return link_pred_loss.mean()
  
  
  class GraphSAGE(nn.Module):
      r"""GraphSAGE from `"Inductive Representation Learning on Large Graphs" <https://arxiv.org/pdf/1706.02216.pdf>`__.
  
      ..math::
          h^{i+1}_{\mathcal{N}(v)}=AGGREGATE_{k}(h_{u}^{k})
          h^{k+1}_{v} = \sigma(\mathbf{W}^{k}·CONCAT(h_{v}^{k}, h_{\mathcal{N}(v)}))
  
      Args:
          in_feats (int) : Size of each input sample.
          hidden_dim (int) : Size of hidden layer dimension.
          out_feats (int) : Size of each output sample.
          num_layers (int) : Number of GraphSAGE Layers.
          dropout (float, optional) : Size of dropout, default: ``0.5``.
          normalize (bool, optional) : Normalze features after each layer if True, default: ``True``.
      """
  
      def __init__(
          self, in_feats, hidden_dim, out_feats, num_layers, dropout=0.5, normalize=False, concat=False, use_bn=False
      ):
          super(GraphSAGE, self).__init__()
          self.convlist = nn.ModuleList()
          self.bn_list = nn.ModuleList()
          self.num_layers = num_layers
          self.dropout = dropout
          self.use_bn = use_bn
          aggr = "concat" if concat else "mean"
          if num_layers == 1:
              self.convlist.append(SAGELayer(in_feats, out_feats, normalize, aggr))
          else:
              self.convlist.append(SAGELayer(in_feats, hidden_dim, normalize, aggr))
              if use_bn:
                  self.bn_list.append(nn.BatchNorm1d(hidden_dim))
              for _ in range(num_layers - 2):
                  self.convlist.append(SAGELayer(hidden_dim, hidden_dim, normalize, aggr))
                  if use_bn:
                      self.bn_list.append(nn.BatchNorm1d(hidden_dim))
              self.convlist.append(SAGELayer(hidden_dim, out_feats, normalize, aggr))
  
      def forward(self, graph, x):
          h = x
          for i in range(self.num_layers - 1):
              h = F.dropout(h, p=self.dropout, training=self.training)
              h = self.convlist[i](graph, h)
              if self.use_bn:
                  h = self.bn_list[i](h)
          return self.convlist[self.num_layers - 1](graph, h)
  
  
  class BatchedGraphSAGE(nn.Module):
      r"""GraphSAGE with mini-batch
  
      Args:
          in_feats (int) : Size of each input sample.
          out_feats (int) : Size of each output sample.
          use_bn (bool) : Apply batch normalization if True, default: ``True``.
          self_loop (bool) : Add self loop if True, default: ``True``.
      """
  
      def __init__(self, in_feats, out_feats, use_bn=True, self_loop=True):
          super(BatchedGraphSAGE, self).__init__()
          self.self_loop = self_loop
          self.use_bn = use_bn
          self.weight = nn.Linear(in_feats, out_feats, bias=True)
  
          nn.init.xavier_uniform_(self.weight.weight.data, gain=nn.init.calculate_gain("relu"))
  
      def forward(self, x, adj):
          device = x.device
          if self.self_loop:
              adj = adj + torch.eye(x.shape[1]).to(device)
          adj = adj / adj.sum(dim=1, keepdim=True)
          h = torch.matmul(adj, x)
          h = self.weight(h)
          h = F.normalize(h, dim=2, p=2)
          h = F.relu(h)
          # TODO: shape = [a, 0, b]
          # if self.use_bn and h.shape[1] > 0:
          #     self.bn = nn.BatchNorm1d(h.shape[1]).to(device)
          #     h = self.bn(h)
          return h
  
  
  class BatchedDiffPoolLayer(nn.Module):
      r"""DIFFPOOL from paper `"Hierarchical Graph Representation Learning
      with Differentiable Pooling" <https://arxiv.org/pdf/1806.08804.pdf>`__.
  
      .. math::
          X^{(l+1)} = S^{l)}^T Z^{(l)}
          A^{(l+1)} = S^{(l)}^T A^{(l)} S^{(l)}
          Z^{(l)} = GNN_{l, embed}(A^{(l)}, X^{(l)})
          S^{(l)} = softmax(GNN_{l,pool}(A^{(l)}, X^{(l)}))
  
      Parameters
      ----------
      in_feats : int
          Size of each input sample.
      out_feats : int
          Size of each output sample.
      assign_dim : int
          Size of next adjacency matrix.
      batch_size : int
          Size of each mini-batch.
      dropout : float, optional
          Size of dropout, default: ``0.5``.
      link_pred_loss : bool, optional
          Use link prediction loss if True, default: ``True``.
      """
  
      def __init__(
          self, in_feats, out_feats, assign_dim, batch_size, dropout=0.5, link_pred_loss=True, entropy_loss=True
      ):
          super(BatchedDiffPoolLayer, self).__init__()
          self.assign_dim = assign_dim
          self.dropout = dropout
          self.use_link_pred = link_pred_loss
          self.batch_size = batch_size
          self.embd_gnn = SAGELayer(in_feats, out_feats, normalize=False)
          self.pool_gnn = SAGELayer(in_feats, assign_dim, normalize=False)
  
          self.loss_dict = dict()
  
      def forward(self, graph, x, batch):
          embed = self.embd_gnn(graph, x)
          pooled = F.softmax(self.pool_gnn(graph, x), dim=-1)
          device = x.device
          masked_tensor = []
          value_set, value_counts = torch.unique(batch, return_counts=True)
          batch_size = len(value_set)
          for i in value_counts:
              masked = torch.ones((i, int(pooled.size()[1] / batch_size)))
              masked_tensor.append(masked)
          masked = torch.FloatTensor(block_diag(*masked_tensor)).to(device)
  
          result = torch.nn.functional.softmax(masked * pooled, dim=-1)
          result = result * masked
          result = result / (result.sum(dim=-1, keepdim=True) + 1e-13)
          # result = masked_softmax(pooled, masked, memory_efficient=False)
  
          h = torch.matmul(result.t(), embed)
          adj = torch.sparse_coo_tensor(torch.stack(graph.edge_index), graph.edge_weight)
          adj_new = torch.sparse.mm(adj, result)
          adj_new = torch.mm(result.t(), adj_new)
  
          if self.use_link_pred:
              adj_loss = torch.norm((adj.to_dense() - torch.mm(result, result.t()))) / np.power((len(batch)), 2)
              self.loss_dict["adj_loss"] = adj_loss
          entropy_loss = (torch.distributions.Categorical(probs=pooled).entropy()).mean()
          assert not torch.isnan(entropy_loss)
          self.loss_dict["entropy_loss"] = entropy_loss
          return adj_new, h
  
      def get_loss(self):
          loss_n = 0
          for _, value in self.loss_dict.items():
              loss_n += value
          return loss_n
  
  
  class BatchedDiffPool(nn.Module):
      r"""DIFFPOOL layer with batch forward
  
      Parameters
      ----------
      in_feats : int
          Size of each input sample.
      next_size : int
          Size of next adjacency matrix.
      emb_size : int
          Dimension of next node feature matrix.
      use_bn : bool, optional
          Apply batch normalization if True, default: ``True``.
      self_loop : bool, optional
          Add self loop if True, default: ``True``.
      use_link_loss : bool, optional
          Use link prediction loss if True, default: ``True``.
      use_entropy : bool, optioinal
          Use entropy prediction loss if True, default: ``True``.
      """
  
      def __init__(
          self, in_feats, next_size, emb_size, use_bn=True, self_loop=True, use_link_loss=False, use_entropy=True
      ):
          super(BatchedDiffPool, self).__init__()
          self.use_link_loss = use_link_loss
          self.use_bn = use_bn
          self.feat_trans = BatchedGraphSAGE(in_feats, emb_size)
          self.assign_trans = BatchedGraphSAGE(in_feats, next_size)
  
          self.link_loss = LinkPredLoss()
          self.entropy = EntropyLoss()
  
          self.loss_module = nn.ModuleList()
          if use_link_loss:
              self.loss_module.append(LinkPredLoss())
          if use_entropy:
              self.loss_module.append(EntropyLoss())
          self.loss = {}
  
      def forward(self, x, adj):
          h = self.feat_trans(x, adj)
          next_l = F.softmax(self.assign_trans(x, adj), dim=-1)
  
          h = torch.matmul(next_l.transpose(-1, -2), h)
          next = torch.matmul(next_l.transpose(-1, -2), torch.matmul(adj, next_l))
  
          for layer in self.loss_module:
              self.loss[str(type(layer).__name__)] = layer(adj, next, next_l)
  
          return h, next
  
      def get_loss(self):
          value = 0
          for _, v in self.loss.items():
              value += v
          return value
  
  
  def toBatchedGraph(batch_adj, batch_feat, node_per_pool_graph):
      adj_list = [
          batch_adj[i : i + node_per_pool_graph, i : i + node_per_pool_graph]
          for i in range(0, batch_adj.size()[0], node_per_pool_graph)
      ]
      feat_list = [batch_feat[i : i + node_per_pool_graph, :] for i in range(0, batch_adj.size()[0], node_per_pool_graph)]
      adj_list = list(map(lambda x: torch.unsqueeze(x, 0), adj_list))
      feat_list = list(map(lambda x: torch.unsqueeze(x, 0), feat_list))
      adj = torch.cat(adj_list, dim=0)
      feat = torch.cat(feat_list, dim=0)
      return adj, feat
  
  
  class DiffPool(BaseModel):
      r"""DIFFPOOL from paper `Hierarchical Graph Representation Learning
      with Differentiable Pooling <https://arxiv.org/pdf/1806.08804.pdf>`__.
  
      Parameters
      ----------
      in_feats : int
          Size of each input sample.
      hidden_dim : int
          Size of hidden layer dimension of GNN.
      embed_dim : int
          Size of embeded node feature, output size of GNN.
      num_classes : int
          Number of target classes.
      num_layers : int
          Number of GNN layers.
      num_pool_layers : int
          Number of pooling.
      assign_dim : int
          Embedding size after the first pooling.
      pooling_ratio : float
          Size of each poolling ratio.
      batch_size : int
          Size of each mini-batch.
      dropout : float, optional
          Size of dropout, default: `0.5`.
      no_link_pred : bool, optional
          If True, use link prediction loss, default: `True`.
      """
  
      @staticmethod
      def add_args(parser):
          parser.add_argument("--num-layers", type=int, default=2)
          parser.add_argument("--num-pooling-layers", type=int, default=1)
          parser.add_argument("--no-link-pred", dest="no_link_pred", action="store_true")
          parser.add_argument("--pooling-ratio", type=float, default=0.15)
          parser.add_argument("--embedding-dim", type=int, default=64)
          parser.add_argument("--hidden-size", type=int, default=64)
          parser.add_argument("--dropout", type=float, default=0.1)
          parser.add_argument("--batch-size", type=int, default=20)
          parser.add_argument("--train-ratio", type=float, default=0.7)
          parser.add_argument("--test-ratio", type=float, default=0.1)
          parser.add_argument("--lr", type=float, default=0.001)
  
      @classmethod
      def build_model_from_args(cls, args):
          return cls(
              args.num_features,
              args.hidden_size,
              args.embedding_dim,
              args.num_classes,
              args.num_layers,
              args.num_pooling_layers,
              int(args.max_graph_size * args.pooling_ratio) * args.batch_size,
              args.pooling_ratio,
              args.batch_size,
              args.dropout,
              args.no_link_pred,
          )
  
      @classmethod
      def split_dataset(cls, dataset, args):
          return split_dataset_general(dataset, args)
  
      def __init__(
          self,
          in_feats,
          hidden_dim,
          embed_dim,
          num_classes,
          num_layers,
          num_pool_layers,
          assign_dim,
          pooling_ratio,
          batch_size,
          dropout=0.5,
          no_link_pred=True,
          concat=False,
          use_bn=False,
      ):
          super(DiffPool, self).__init__()
          self.assign_dim = assign_dim
          self.assign_dim_list = [assign_dim]
          self.use_bn = use_bn
          self.dropout = dropout
          self.use_link_loss = not no_link_pred
          # assert num_layers > 3, "layers > 3"
          self.diffpool_layers = nn.ModuleList()
          self.before_pooling = GraphSAGE(
              in_feats, hidden_dim, embed_dim, num_layers=num_layers, dropout=dropout, use_bn=self.use_bn
          )
          self.init_diffpool = BatchedDiffPoolLayer(
              embed_dim, hidden_dim, assign_dim, batch_size, dropout, self.use_link_loss
          )
  
          pooled_emb_dim = embed_dim
          self.after_pool = nn.ModuleList()
          after_per_pool = nn.ModuleList()
          for _ in range(num_layers - 1):
              after_per_pool.append(BatchedGraphSAGE(hidden_dim, hidden_dim))
          after_per_pool.append(BatchedGraphSAGE(hidden_dim, pooled_emb_dim))
          self.after_pool.append(after_per_pool)
  
          for _ in range(num_pool_layers - 1):
              self.assign_dim = int(self.assign_dim // batch_size * pooling_ratio) * batch_size
              self.diffpool_layers.append(
                  BatchedDiffPool(
                      pooled_emb_dim, self.assign_dim, hidden_dim, use_bn=self.use_bn, use_link_loss=self.use_link_loss
                  )
              )
  
              for _ in range(num_layers - 1):
                  after_per_pool.append(BatchedGraphSAGE(hidden_dim, hidden_dim))
              after_per_pool.append(BatchedGraphSAGE(hidden_dim, pooled_emb_dim))
              self.after_pool.append(after_per_pool)
  
              self.assign_dim_list.append(self.assign_dim)
  
          if concat:
              out_dim = pooled_emb_dim * (num_pool_layers + 1)
          else:
              out_dim = pooled_emb_dim
          self.fc = nn.Linear(out_dim, num_classes)
  
      def reset_parameters(self):
          for i in self.modules():
              if isinstance(i, nn.Linear):
                  nn.init.xavier_uniform_(i.weight.data, gain=nn.init.calculate_gain("relu"))
                  if i.bias is not None:
                      nn.init.constant_(i.bias.data, 0.0)
  
      def after_pooling_forward(self, gnn_layers, adj, x, concat=False):
          readouts = []
          h = x
          for layer in gnn_layers:
              h = layer(h, adj)
              readouts.append(h)
          # readout = torch.cat(readouts, dim=1)
          return h
  
      def forward(self, batch):
          readouts_all = []
  
          init_emb = self.before_pooling(batch, batch.x)
          adj, h = self.init_diffpool(batch, init_emb, batch.batch)
          value_set, value_counts = torch.unique(batch.batch, return_counts=True)
          batch_size = len(value_set)
          adj, h = toBatchedGraph(adj, h, adj.size(0) // batch_size)
          h = self.after_pooling_forward(self.after_pool[0], adj, h)
          readout = torch.sum(h, dim=1)
          readouts_all.append(readout)
  
          for i, diff_layer in enumerate(self.diffpool_layers):
              h, adj = diff_layer(h, adj)
              h = self.after_pooling_forward(self.after_pool[i + 1], adj, h)
              readout = torch.sum(h, dim=1)
              readouts_all.append(readout)
          pred = self.fc(readout)
          return pred
  
      def graph_classificatoin_loss(self, batch):
          pred = self.forward(batch)
          pred = F.log_softmax(pred, dim=-1)
          loss_n = F.nll_loss(pred, batch.y)
          loss_n += self.init_diffpool.get_loss()
          for layer in self.diffpool_layers:
              loss_n += layer.get_loss()
          return loss_n
  
apis: torch.ones_like torch.zeros_like torch.nn.Embedding torch.nn.CrossEntropyLoss torch.nn.Dropout
apis: numpy.sum torch.nn.functional.log_softmax torch.nn.functional.nll_loss sklearn.metrics.f1_score torch.is_tensor torch.nn.BCEWithLogitsLoss torch.cat torch.mean
apis: torch.stack torch.nn.functional.softmax torch.nn.functional.nll_loss torch.nn.ModuleList torch.eye torch.cat torch.nn.functional.dropout torch.nn.BatchNorm1d torch.unique torch.sparse.mm torch.unsqueeze torch.distributions.Categorical torch.nn.functional.normalize scipy.linalg.block_diag torch.isnan torch.sum torch.nn.init.calculate_gain torch.nn.functional.log_softmax torch.nn.Linear torch.nn.init.constant_ torch.nn.functional.relu torch.matmul