DiffPool with nn layers

While there is an example of the DiffPool paper, it is not clear how to use with existing nn layers like GraphConv or RelGraphConv. The example uses a custom SageGraph model which explicitly takes the features and adgacency matrix as inputs and goes on to further so the message passing.

The diffPooling operation generates an updated feature matrix and an updated adgacency matrix.
With the current nn modules is there a way of sending the updated adgacency matrix for a batched graph input such that the message update and apply funcrion work on basis of this new structural information.

Sorry for the late reply. It’s a bit tricky because the graph after diffpool is not an actually graph. The new adjacency matrix is a real matrix instead of a 0-1 matrix. Therefore the graph structure is hard to determine on that.