Flax
Neural network library for JAX by Google providing Linen module system.
About
Flax is a neural network library for JAX from Google, originally built by Google Brain with the JAX team. It provides module systems for defining models, training-state management, and serialization, and is used across many Google research projects. The newer Flax NNX API adds Python reference semantics so models behave like regular mutable objects, while the earlier Linen API remains supported. Released under the Apache 2.0 license.
Reviews (0)
Leave a Review
No reviews yet. Be the first to review!
Details
- Category
- AI Frameworks & Libraries
- Price
- Free
- Platform
- Local/Desktop
- Difficulty
- Advanced (4/5)
- License
- Apache-2.0
- Added
- Apr 3, 2026
Related Tools
JAX neural network library with PyTorch-like module system.
High-level deep learning API supporting JAX, TensorFlow, and PyTorch backends.
Deep learning platform by Baidu with industrial-strength deployment tools.
Hugging Face library for state-of-the-art NLP models
Deploy LangChain runnables as REST APIs
End-to-end open-source ML platform by Google for training and deployment.