Optax

E438357

Optax is a gradient processing and optimization library for JAX, providing a flexible collection of composable optimizers and transformations for training machine learning models.

All labels observed (1)

Label Occurrences
Optax canonical 1

How this entity was disambiguated

Statements (49)

Predicate Object
instanceOf JAX ecosystem library
machine learning library
open-source software
software library
compatibleWith jax.grad
jax.jit
jax.pmap
jax.vmap
designGoal compatibility with JAX transformations
composability of optimization components
functional programming style API
separation of model definition and optimization logic
documentationFormat examples and tutorials
online API reference
ecosystem DeepMind JAX ecosystem NERFINISHED
feature adaptive gradient methods
composable optimizers
gradient clipping
gradient transformations
learning rate schedules
loss scaling utilities
momentum methods
multi-device training utilities
parameter regularization
second-order optimization utilities
stateful and stateless transformations
stochastic gradient descent variants
isFreeAndOpenSource true
license Apache License 2.0
primaryDomain gradient-based optimization
optimization for machine learning
programmingLanguage Python
providesOptimizer adafactor NERFINISHED
adagrad
adam NERFINISHED
adamw
lamb
lars
novograd NERFINISHED
rmsprop NERFINISHED
sgd
yogi
repositoryPlatform GitHub NERFINISHED
supportsFramework JAX NERFINISHED
supportsUseCase large-scale distributed training
reinforcement learning optimization
research prototyping in optimization
training neural networks
typicalWorkflow init-update functional optimizer API

How these facts were elicited

Referenced by (1)

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

JAX compatibleWith Optax