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.
A 23-video series teaching JAX is also available
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!