orbax (pypi): Orbax
Keywords: JAX|machine|learning