I have an existing implementation (written by somebody else) of an MPNN using the graph_nets library. The graph net is based on a tree, but has 4 times as many edges: if U is the parent of V and R is the root of the tree, then the graph net has the edges (U, V), (V, U), (R, V) and (V, R).
I think it makes sense to have a separate MLP for each kind of edge. So I would have for each edge a single feature, whose value is an integer from 0 to 3, and the edge function would select one of 4 MLPs and apply it.
Does this make sense? Is there a way to implement this with graph_nets? How?
