jax-metrics (pypi): A metrics library for the JAX ecosystem
Keywords: jax|python