Flax documentationΒΆ

Flax is a neural network library and ecosystem for JAX that is designed for flexibility. Flax is in use by a growing community of researchers and engineers at Google who happily use Flax for their daily research.

For a quick introduction and short example snippets, see our README.