Is there a way to have multiple node types and edge types so that the update and apply functions use a different neural network (and/or different parameters) for each type?
Hi, you may want to take a look at our tutorial on RGCN where different weights are used for different edge types. In simple cases where we are only performing matrix multiplications with different weights for different edge types, we can simply initialize the weight as a tensor of shape
(num_edge_types, in_size, out_size). Then in message passing we can index the weight matrix using edge types and perform batched matrix multiplication at once.
In more complex cases as long as the edge types are fixed, we can create a mapping from edge types to edges at the beginning and process all edges of the same type at a time. Hope this helps.
I’ll check that out! Thanks!