clu (pypi): Set of libraries for ML training loops in JAX.
Keywords: JAX|machine|learning|flax|jax|ml