Flax

Training Models at Scale Tutorials

Deep Learning Tutorials on Training Models at Scale in JAX with Flax

GDE: Machine Learning with JAX and Flax

Introduction talk for learning JAX and Flax as Deep Learning framework

Deep Learning Tutorials Translated to JAX with Flax

Deep Learning Tutorials in PyTorch with PyTorch Lightning and JAX with Flax