Writing a Training Loop in JAX + FLAX by Saurav Maheshkar and Soumik Rakshit
A tutorial on writing a simple end-to-end training and evaluation pipeline in JAX, Flax and Optax.
- from
- JAX
- added
- 2026-10-10
- likes
- 0
similar
-
Deep Learning tutorials with JAX+Flax by Phillip Lippe uvadlc-notebooks.readthedocs.io
A series of notebooks explaining various deep learning concepts, from basics (e.g. intro to JAX/Flax, activiation functions) to recent advances (e.g., Vision Transformers, SimCLR), with translations to PyTorch.
-
Meta-Learning in 50 Lines of JAX by Eric Jang blog.evjang.com
Introduction to both JAX and Meta-Learning.
-
-
EasyLM github.com
LLMs made easy: Pre-training, finetuning, evaluating and serving LLMs in JAX/Flax.
-
-
Accurate Quantized Training github.com
Tools and libraries for running and analyzing neural network quantization experiments in JAX and Flax.
JAX › Tutorials and Blog Posts: “A tutorial on writing a simple end-to-end training and evaluation pipeline in JAX, Flax and Optax.”