Support for GNNs

[ataa1312]

Hello Flax team,

I’ve been exploring graph neural network support in the JAX ecosystem and noticed that libraries like Jraph and a few others haven’t seen updates in nearly three years. Are there any plans to introduce native GNN modules in Flax, or could you recommend a well-maintained, PyG-like graph library for JAX?

Thank you for your time and all the great work on Flax!