google/flax: Flax is a neural network library for JAX that is designed for flexibility.
Category: ml-dl
Language: Python
Topics: jax, neural-network
Homepage: https://flax.readthedocs.io