jax-loss (pypi): A jax library of common machine learning loss functions
Keywords: jax|loss-functions|machine-learning