Mol line graph in jtnn example

I’m trying to implement neural message passing for quantum chemistry ( in dgl and I’m following the message passing model example in JTNN but I’m confused about mol_line_graph update method at I’m wondering why do we update the line graph in the for loop with depth and do update_all for mol_graph just once.

Thanks in advance!

The implementation follows Junction Tree Variational Autoencoder:

The for loop corresponds to the loopy-BP-like message passing in Equation 1 (which happens on edges, or equivalently, nodes of the line graph).

The last update_all simply gathers the edge states to nodes as in Equation 2.