optax
Optax is a gradient processing and optimization library for JAX.
2k stars
34 watching
179 forks
Language: Python
last commit: 4 days ago
Linked from 2 awesome lists
machine-learningoptimization
optax Optax is a gradient processing and optimization library for JAX.
krzjoa/awesome-python-data-science
n2cholas/awesome-jax