Flax is a neural network ecosystem for JAX that is designed for flexibility.
By chatting or signing in you agree to the Terms and chat-message logging (revocable in History).