GitHub repository: stefanocortinovis/optax
Primary language: Python
Description: Optax is a gradient processing and optimization library for JAX.