What is the output of a TopKPooling layer?

I am trying to implement a TopKPooling layer but the output is not documented. I am get a tuple of 6 objects – the first two items look to be the new node feature matrix and edge indices, the third is a None type (probably because I am doing standard TopKPooling), the fourth looks to be the new batch tensor which indicates which of the nodes belong to which graph. The fifth and sixth I am less sure on: they are both 1D tensors of length |V’| where V’ is the new coarsened node set – the fifth looks to be strictly integers with no gradient attached so it could be some kind of labelling and the sixth looks to be floats with a gradient.

could some please help me understand what these outputs are, thanks :slight_smile:

The code I am using is as follows:

class KPooling(nn.Module):

    def __init__(self, ratio, input_size, hidden_size, output_size):

        super(KPooling, self).__init__()

        self.gcn1 = SAGEConv(input_size, hidden_size)
        self.first_pool_layer = TopKPooling(hidden_size, ratio)

    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch

        x = self.gcn1(x, edge_index)
        x, edge_index, _, batch, ?, ? = self.first_pool_layer(x, edge_index, batch=batch)

where the ? I’ve added to show that I am unsure on what these two tensor outputs are.