flax

Anaconda Verified

Flax A neural network library for JAX designed for flexibility

Versions

Installation

To install this package, run one of the following:

Conda

$conda install main::flax

Usage Tracking

Version

1 / 8 versions selected

Downloads (Last 6 months): 0

Description

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

About

Summary

Flax A neural network library for JAX designed for flexibility

Last Updated

Feb 24, 2026 at 14:23

License

Apache-2.0

Total Downloads

Version Downloads

Supported Platforms

linux-aarch64

macOS-arm64

linux-64

Home

https://github.com/google/flax

GitHub Repository

https://github.com/google/flax

Documentation

https://flax.readthedocs.io