JAX-Privacy: Algorithms for Privacy-Preserving Machine Learning in JAX

| Docs | Overview | Paper | Citing | Contact

JAX-Privacy is a library designed to simplify the deployment of robust and performant mechanisms for differentially private (DP) machine learning in JAX. Guided by design principles of usability, flexibility, and efficiency, it provides modular, verified primitives alongside high-level training APIs for both researchers requiring deep customization and practitioners seeking an out-of-the-box experience.

For installation instructions, examples, and full API documentation, please visit the JAX Privacy documentation or read the accompanying paper.

Scope and Design Philosophy

JAX Privacy is a general-purpose Python library for differentially private machine learning, with a focus on DP-SGD-style mechanisms. Rather than implementing task-specific end-to-end pipelines, it provides both high-level training APIs ([training.py][training-py] and the [Keras API][keras-doc]) and lower-level, composable building blocks that integrate into any JAX training loop. For a detailed walkthrough of the five core building blocks and the three API tiers, see the [Library Overview][overview-doc].

What Is in Scope

• Composable DP-SGD Ingredients: Pure Python/JAX modules for per-example

• Modality- and Model-Agnostic Computation: JAX Privacy works with any

• Zero Framework Lock-In: JAX Privacy intentionally does not bring in any

What Is Out of Scope

• Models and Data Loaders: Defining model architectures, tokenizers, and

• Domain-Specific End-to-End Pipelines: Task-specific workflows such as

• Binaries and Surrounding Infrastructure: As a library, JAX Privacy

[overview-doc]: https://jax-privacy.readthedocs.io/en/latest/overview.html [keras-doc]: https://jax-privacy.readthedocs.io/en/latest/keras_api.html [training-py]: https://github.com/google-parfait/jax_privacy/blob/main/jax_privacy/training.py [examples-dir]: https://github.com/google-parfait/jax_privacy/tree/main/examples

How to Cite This Repository

If you use JAX-Privacy in your work, please cite the accompanying paper:

@article{mckenna2026jaxprivacy,
  author = {McKenna, Ryan and Andrew, Galen and Balle, Borja and
Doroshenko, Vadym and Ganesh, Arun and Kong, Weiwei and Kurakin, Alex and
McMahan, Brendan and Pravilov, Mikhail},
  title = {{JAX}-{P}rivacy: A library for differentially private machine learning},
  journal = {arXiv preprint arXiv:2602.17861},
  url = {https://arxiv.org/abs/2602.17861},
  year = {2026},
}

To cite the software repository directly:

@software{jax-privacy2022github,
  author = {Balle, Borja and Berrada, Leonard and Charles, Zachary and
Choquette-Choo, Christopher A and De, Soham and Doroshenko, Vadym and Dvijotham,
Dj and Galen, Andrew and Ganesh, Arun and Ghalebikesabi, Sahra and Hayes, Jamie
and Kairouz, Peter and McKenna, Ryan and McMahan, Brendan and Pappu, Aneesh and
Ponomareva, Natalia and Pravilov, Mikhail and Rush, Keith and Smith, Samuel L
and Stanforth, Robert and Mishra, Chaitanya},
  title = {{JAX}-{P}rivacy: Algorithms for Privacy-Preserving Machine Learning in JAX},
  url = {http://github.com/google-parfait/jax_privacy},
  version = {2.3.0.dev0},
  year = {2026},
}

Contact

If you have any questions or feedback, you can contact us via email: [email protected].

Acknowledgements

• NFNet codebase
• DeepMind JAX Ecosystem

License

All code is made available under the Apache 2.0 License. Model parameters are made available under the Creative Commons Attribution 4.0 International (CC BY 4.0) License.

See https://creativecommons.org/licenses/by/4.0/legalcode for more details.

Disclaimer

This is not an official Google product.