What's the difference between graph_nets and jraph? about graph_nets HOT 3 CLOSED

I also read some demo code from jraph and it seems that the model there obay the Alogorith1 from https://arxiv.org/abs/1806.01261.

from graph_nets.

Thanks for the comment:

What's the difference between graph_nets and jraph?
The main difference is that jraph works with JAX, and graph_nets works with TensorFlow. We also provided an improved API for jraph which offers native support for nested node, edge and global features.

1, Could I use graph_nets to train a GNN model, such as GCN/GraphSage/GAT ?
Yes, although in principle you would have to write the model you want based on the building blocks. For example, for a GCN, you can simply broadcast senders nodes to edges, and then aggregate to receivers. And for GATs, you may be able to reuse the (or parts of) the self attention module: https://github.com/deepmind/graph_nets/blob/64771dff0d74ca8e77b1f1dcd5a7d26634356d61/graph_nets/modules.py#L657

it seems that graph_nets is used to define graph-structured data and the GNN model is defined and trained by jraph, is that right?

In that project we used TensorFlow for the data pipelines (to be able to benefit from fast prefetching), and JAX for the model, hence the use of the two graph nets libraries.

Hope this helps!

from graph_nets.

great thx@ @alvarosg

from graph_nets.