def METHOD_NAME(
    self,
    x: Tensor,
    edge_index: Tensor,
    edge_attr: OptTensor = None,
    batch: OptTensor = None,
    attn: OptTensor = None,
) -> Tuple[Tensor, Tensor, OptTensor, OptTensor, Tensor, Tensor]:
    r"""
    Args:
        x (torch.Tensor): The node feature matrix.
        edge_index (torch.Tensor): The edge indices.
        edge_attr (torch.Tensor, optional): The edge features.
            (default: :obj:`None`)
        batch (torch.Tensor, optional): The batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns
            each node to a specific example. (default: :obj:`None`)
        attn (torch.Tensor, optional): Optional node-level matrix to use
            for computing attention scores instead of using the node
            feature matrix :obj:`x`. (default: :obj:`None`)
    """
    if batch is None:
        batch = edge_index.new_zeros(x.size(0))
    attn = x if attn is None else attn
    attn = attn.view(-1, 1) if attn.dim() == 1 else attn
    attn = self.gnn(attn, edge_index)
    select_out = self.select(attn, batch)
    perm = select_out.node_index
    score = select_out.weight
    assert score is not None
    x = x[perm] * score.view(-1, 1)
    x = self.multiplier * x if self.multiplier != 1 else x
    connect_out = self.connect(select_out, edge_index, edge_attr, batch)
    return (x, connect_out.edge_index, connect_out.edge_attr,
            connect_out.batch, perm, score)