External Resources

JAX and Graph Learning Resources

• JAX Documentation: Official JAX documentation [Website]

• JAX Tutorials: Comprehensive JAX learning materials [Website]

• Flax NNX Guide: Modern neural networks with Flax [Documentation]

Graph Machine Learning

• Stanford CS224W: Machine Learning with Graphs lectures [YouTube]

• Graph Neural Networks: Comprehensive introduction [Distill.pub]

• Geometric Deep Learning: Theoretical foundations [Website]

PyTorch Geometric Integration

Since JraphX can work with PyTorch Geometric datasets:

• PyTorch Geometric: Original library for graph neural networks [Website, GitHub]

• PyG Datasets: Comprehensive collection of graph datasets [Documentation]

• Graph learning datasets: Curated collection [TUDatasets]

JAX Ecosystem

• JAX AI Stack [GitHub]

Related Libraries

• Jraph: DeepMind’s archived graph neural network library [GitHub (archived)]

• JAX-MD: Molecular dynamics with JAX [GitHub]

• Graph Nets: DeepMind’s graph networks library [GitHub] in TensorFlow/Sonnet

Contributing

JraphX is an open-source project. Contributions are welcome!

• Issues: Report bugs or request features

• Pull Requests: Contribute code improvements

• Documentation: Help improve these docs

• Examples: Share your JraphX applications