Last released Mar 14, 2024
Flax: A neural network library for JAX designed for flexibility
Supported by