Learning-JAX

Slide decks, coding exercises, and quick references for learning the JAX AI Stack. The coding exercises are designed to be runnable in a free Colab instance.

For more comprehensive documentation please see the individual websites:

• https://jaxstack.ai
• https://jax.dev
• https://flax.readthedocs.io
• https://orbax.readthedocs.io
• https://optax.readthedocs.io
• https://google-grain.readthedocs.io
• https://chex.readthedocs.io

Join our growing community on Discord and connect with other developers!