Documentation
API
Examples
About
VAE example with Haiku and JAXopt.
Image classification example with Haiku and JAXopt.
Comparison of different SGD algorithms.
SPMD ResNet example with Flax and JAXopt.