More losses!

[h-0-0]

Hi, I thought it could be nice if there were a few more losses ready to go in mlx. To start with I was thinking:

• L1_loss (L1 loss)
• mse_loss (mean squared error)

Maybe also:

• nll_loss (negative log likelihood loss)
• kl_div_loss (KL-divergence loss)

Think this should pretty straightforward to do? Simply a case of adding them to python/mlx/nn/losses.py

Any thoughts? And I am happy to work on this issue.

[awni]

Definitely! L1 and MSE for sure. I would also add a binary cross entropy as well.

We already started a loss module here but it has no docs yet.

Some testing + docs for any additional losses would be fantastic.

Also, what's the distinction between the nll_loss and the cross_entropy loss we already have?

[h-0-0]

Hi, I'll implement the L1 and MSE losses in that case, and you are correct about the negative log likelihood! I was just going through the pytorch losses and thinking about which would be the most helpful to add. Of course I'll add docs and testing aswell.

[awni]

There have already been contributions for L1 and MSE. You can see the referenced PRs in this issue to see what's been/being added.

If there are other common losses you'd like to add, feel free to let us know and we can discuss here.

[awni]

All the losses in this issue have been pushed to main so I will close it. Feel free to open a new issue for new losses.