optax (pypi): A gradient processing and optimisation library in JAX.
Keywords: reinforcement-learning|python|machine|learning