Flax

E438356

Flax is a neural network library for JAX that provides a flexible, modular framework for building and training machine learning models in Python.

All labels observed (1)

Label Occurrences
Flax canonical 1

How this entity was disambiguated

Statements (48)

Predicate Object
instanceOf neural network library
open-source software
software library
basedOn JAX NERFINISHED
compatibleWith NumPy API via JAX NERFINISHED
Optax NERFINISHED
Orbax NERFINISHED
designGoal composability
explicit handling of state and randomness
flexibility
developedBy Google NERFINISHED
Google Research NERFINISHED
domain deep learning
machine learning
ecosystem JAX ecosystem
hasComponent flax.linen
flax.optim (deprecated in favor of Optax) NERFINISHED
flax.training
hasFeature checkpointing utilities
module abstraction (flax.linen)
parameter management
random number handling
serialization utilities
state management
training loops utilities
hostedOn GitHub NERFINISHED
license Apache License 2.0
programmingLanguage Python
provides flexible neural network framework
modular framework for building models
tools for training machine learning models
repositoryURL https://github.com/google/flax
supports JAX transformations
NLP models
convolutional neural networks
functional programming style
gradient-based optimization
model training
neural network construction
recurrent neural networks
reinforcement learning models
stateful neural network modules
transformer models
vision models
supportsLanguage Python
usedFor production ML systems
rapid prototyping of models
research in machine learning

How these facts were elicited

Referenced by (1)

Full triples — surface form annotated when it differs from this entity's canonical label.

JAX compatibleWith Flax