Deep learning

VAE example with Haiku and JAXopt.

VAE example with Haiku and JAXopt.

Image classification example with Haiku and JAXopt.

Image classification example with Haiku and JAXopt.

Comparison of different SGD algorithms.

Comparison of different SGD algorithms.

SPMD ResNet example with Flax and JAXopt.

SPMD ResNet example with Flax and JAXopt.