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