Flax A neural network library for JAX designed for flexibility
conda install main::flax
Flax provides a flexible end-to-end user experience for researchers and developers who use JAX for neural networks. Flax enables you to use the full power of JAX
