def forward(
        self,
        graph: DGLGraph,
        h_pos: Tensor,
        h_neg: Tensor,
        bias_pos: Optional[Tensor] = None,
        bias_neg: Optional[Tensor] = None,
    ):
        """
        Description
        -----------
        Perform index selection

        Parameters
        ----------
        graph : dgl.DGLGraph
            Input graph.
        h_pos : torch.Tensor
            The node features of positive samples
            It has the same shape as :obj:`h_x`
        h_neg : torch.Tensor
            The node features of negative samples
            It has the same shape as :obj:`h_x`
        bias_pos : torch.Tensor
            Bias parameter vector for positive scores
            shape: :obj:`(num_nodes)`
        bias_neg : torch.Tensor
            Bias parameter vector for negative scores
            shape: :obj:`(num_nodes)`
        """
        # compute scores
        h_pos = self.dense(h_pos)
        h_neg = self.dense(h_neg)
        embed = self.gcn(graph, h_pos)
        h_center = torch.sigmoid(embed)

        logit, logit_pos = self.discriminator(
            h_center, h_pos, h_neg, bias_pos, bias_neg
        )
        scores = torch.sigmoid(logit_pos)

        # sort scores
        scores, idx = torch.sort(scores, descending=True)

        # select top-k
        num_nodes = graph.num_nodes()
        num_select_nodes = int(self.pool_ratio * num_nodes)
        size_list = [num_select_nodes, num_nodes - num_select_nodes]
        select_scores, _ = torch.split(scores, size_list, dim=0)
        select_idx, non_select_idx = torch.split(idx, size_list, dim=0)

        return logit, select_scores, select_idx, non_select_idx, embed

        Description
        -----------
        Perform index selection

        Parameters
        ----------
        graph : dgl.DGLGraph
            Input graph.
        h_pos : torch.Tensor
            The node features of positive samples
            It has the same shape as :obj:`h_x`
        h_neg : torch.Tensor
            The node features of negative samples
            It has the same shape as :obj:`h_x`
        bias_pos : torch.Tensor
            Bias parameter vector for positive scores
            shape: :obj:`(num_nodes)`
        bias_neg : torch.Tensor
            Bias parameter vector for negative scores
            shape: :obj:`(num_nodes)`
        
examples/pytorch/gxn/layers.py