What is the best way to implement a Tree Convolution?

I was reading the paper Tree Gen (https://arxiv.org/pdf/1911.09983.pdf) and wanted to implement the TreeConvolution:

Screen Shot 2021-07-27 at 9.16.54 AM

I was wondering if this was possible to be done with PyTorch Geometric or DGL or something like that or if the best way was to simply generate the adjacency matrix, multiply my vectors with it and then do a normal convolution?

cross-posted:

Hi @albanD, sorry for the direct post but I wasn’t sure who else to tag. Is there a way to “insert a (any) GraphNN layer” in a feedforward model I already have in pytorch (e.g. pytorch geometric)?

e.g.

imagine I wanted to implement a sequential model with a GCN:

f = nn.Sequential(OrderedDict([
    ('f1', nn.Linear(Din, Dout)),
    ('out', nn.ReLU()),
    ('gcn1', torch_geometric.GCN(*params))
]))
y: Tensor =  f(x)

Something like that. Is it possible?

I am not sure to follow what would prevent you from doing that?
You will need to make sure that input/output shape match. But that’s it no?