ddehueck / jax-loss

A jax library of common machine learning loss functions

Jax Loss

A jax library of common machine learning loss functions

Installation

pip install jax-loss

Loss Functions Included

• Binary Cross Entropy With Logits Loss

If you don't see the loss function you need, feel free to open an issue or open a PR!

Loss Functions to be Implemented

• L1
• Mean Squared Error
• Cross Entropy
• Negative Log Liklihood
• KL Divergance
• Binary Cross Entropy
• Triplet Margin
• Cosine Embedding Loss
• Hinge

jax-loss's Issues

Is bce_w_logits equivalent to the pytorch BCEWithLogitsLoss?

jax-loss/jax_loss/classification.py

Line 4 in 410d1f5

def bce_w_logits(x, y, weight=None, average=True):

The comment seems to indicate that this is indeed true, but I wanted to confirm.

contributing.md

Is this project still live?

Hi, just wondering whether you're still working on this and accepting PRs etc? I'm potentially interested in collaborating and adding some losses. Let me know!

Develop in an example-first manner

We can use vmap for automatic vectorization so we should develop from this example-first framework.

Add tests and badges

Add tests for each loss function - comparing to pytorch and tensorflow implementations.

Could be nice to have badges as described in: https://thomas-cokelaer.info/blog/2014/08/1013/