Need help with awesome-jax?
Click the “chat” button below for chat support from the developer who created it, or find similar developers for support.

About the developer

n2cholas
172 Stars 13 Forks Creative Commons Zero v1.0 Universal 56 Commits 2 Opened issues

Description

JAX - A curated list of resources https://github.com/google/jax

Services available

!
?

Need anything else?

Contributors list

# 229,361
Python
C++
Shell
pytorch
15 commits
# 377,535
Python
normali...
bayesia...
11 commits
# 47,373
Jupyter...
Swift
C++
differe...
2 commits
# 26,241
pandas
Jupyter...
network...
bayesia...
2 commits
# 189,673
The Jul...
Shell
Jupyter...
restric...
1 commit
# 266,655
pytorch
Shell
jax
Swift
1 commit

Awesome JAX AwesomeJAX Logo

JAX brings automatic differentiation and the XLA compiler together through a numpy-like API for high performance machine learning research on accelerators like GPUs and TPUs.

This is a curated list of awesome JAX libraries, projects, and other resources. Contributions are welcome!

Contents

Libraries

  • Neural Network Libraries
    • Flax - Centered on flexibility and clarity.
    • Haiku - Focused on simplicity, created by the authors of Sonnet at DeepMind.
    • Objax - Has an object oriented design similar to PyTorch.
    • Elegy - Implements the Keras API with some improvements.
    • RLax - Library for implementing reinforcement learning agents.
    • Trax - "Batteries included" deep learning library focused on providing solutions for common workloads.
    • Jraph - Lightweight graph neural network library.
    • Neural Tangents - High-level API for specifying neural networks of both finite and infinite width.
  • NumPyro - Probabilistic programming based on the Pyro library.
  • Chex - Utilities to write and test reliable JAX code.
  • Optax - Gradient processing and optimization library.
  • JAX, M.D. - Accelerated, differential molecular dynamics.
  • Coax - Turn RL papers into code, the easy way.
  • SymJAX - Symbolic CPU/GPU/TPU programming.
  • mcx - Express & compile probabilistic programs for performant inference.

New Libraries

This section contains libraries that are well-made and useful, but have not necessarily been battle-tested by a large userbase yet.

  • Neural Network Libraries
    • Parallax - Prototype immutable torch modules for JAX.
    • FedJAX - Federated learning in JAX, built on Optax and Haiku.
  • jax-unirep - Library implementing the UniRep model for protein machine learning applications.
  • jax-flows - Normalizing flows in JAX.
  • sklearn-jax-kernels -
    scikit-learn
    kernel matrices using JAX.
  • jax-cosmo - Differentiable cosmology library.
  • efax - Exponential Families in JAX.
  • mpi4jax - Combine MPI operations with your Jax code on CPUs and GPUs.

Models and Projects

Videos

Papers

This section contains papers focused on JAX (e.g. JAX-based library whitepapers, research on JAX, etc). Papers implemented in JAX are listed in the Models/Projects section.

Blog Posts

Community

Contributing

Contributions welcome! Read the contribution guidelines first.

We use cookies. If you continue to browse the site, you agree to the use of cookies. For more information on our use of cookies please see our Privacy Policy.