I recently made this graph neural networks library. I haven't found anything similar for JAX. There is Jraph, which I like, but I wanted something more straightforward and with more variety of layers/parameters.
Show HN: Haiku Geometric – Graph Neural Networks for JAX | Heykuki News