Citation

If you find JraphX useful for your research, please consider citing it along with the foundational libraries it builds upon.

JraphX

@software{jraphx2025github,
  author = {David Braun},
  title = {{JraphX}: A Graph Neural Network library for {JAX}/{Flax NNX}},
  url = {https://github.com/DBraun/jraphx},
  year = {2025},
}

PyTorch Geometric

JraphX incorporates substantial portions of code and documentation from PyTorch Geometric:

@inproceedings{Fey/Lenssen/2019,
  title={Fast Graph Representation Learning with {PyTorch Geometric}},
  author={Fey, Matthias and Lenssen, Jan E.},
  booktitle={ICLR Workshop on Representation Learning on Graphs and Manifolds},
  year={2019},
}

@inproceedings{Fey/etal/2025,
  title={{PyG} 2.0: Scalable Learning on Real World Graphs},
  author={Fey, Matthias and Sunil, Jinu and Nitta, Akihiro and Puri, Rishi and Shah, Manan, and Stojanovi{\v{c}, Bla{\v{z} and Bendias, Ramona and Alexandria, Barghi and Kocijan, Vid and Zhang, Zecheng and He, Xinwei and Lenssen, Jan E. and Leskovec, Jure},
  booktitle={Temporal Graph Learning Workshop @ KDD},
  year={2025},
}

DeepMind Jraph

JraphX serves as an unofficial successor to DeepMind’s jraph library:

@software{jraph2020github,
  author = {Jonathan Godwin* and Thomas Keck* and Peter Battaglia and Victor Bapst and Thomas Kipf and Yujia Li and Kimberly Stachenfeld and Petar Veli\v{c}kovi\'{c} and Alvaro Sanchez-Gonzalez},
  title = {{J}raph: {A} library for graph neural networks in jax.},
  url = {http://github.com/deepmind/jraph},
  version = {0.0.1.dev},
  year = {2020},
}