zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-18 10:26:56 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Style fix for loss functions (#91)

Browse Source 
* MLE and L1 loss functions

* logsoftmax change and tests

* subtract max logit for numerical stability

* l1 name change

* cross entropy reduction + unit tests

* docstrings

* l1 test name change

* old loss impl + default none

* style
This commit is contained in:
Kai Ma
2023-12-09 00:11:56 -05:00
committed by GitHub
parent 641d316484
commit cb9e585b8e
2 changed files with 15 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
python/mlx/nn/losses.py
View File

@@ -2,7 +2,10 @@
import mlx.core as mx
def cross_entropy(logits: mx.array, targets: mx.array, axis: int = -1, reduction: str = 'none'):
def cross_entropy(
logits: mx.array, targets: mx.array, axis: int = -1, reduction: str = "none"
):
"""
Computes the cross entropy loss between logits and targets.
@@ -10,10 +13,10 @@ def cross_entropy(logits: mx.array, targets: mx.array, axis: int = -1, reduction
logits (mx.array): The predicted logits.
targets (mx.array): The target values.
axis (int, optional): The axis over which to compute softmax. Defaults to -1.
reduction (str, optional): Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'.
'none': no reduction will be applied.
reduction (str, optional): Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'.
'none': no reduction will be applied.
'mean': the sum of the output will be divided by the number of elements in the output.
'sum': the output will be summed.
'sum': the output will be summed.
Defaults to 'none'.
Returns:
@@ -22,15 +25,16 @@ def cross_entropy(logits: mx.array, targets: mx.array, axis: int = -1, reduction
score = mx.take_along_axis(logits, targets[..., None], axis).squeeze(-1)
loss = mx.logsumexp(logits, axis=axis) - score
if reduction == 'mean':
if reduction == "mean":
return mx.mean(loss)
elif reduction == 'sum':
elif reduction == "sum":
return mx.sum(loss)
elif reduction == 'none':
elif reduction == "none":
return loss
else:
raise ValueError("Invalid reduction. Must be 'none', 'mean', or 'sum'.")
def l1_loss(predictions: mx.array, targets: mx.array):
"""
Computes the L1 loss between predictions and targets.
@@ -43,4 +47,3 @@ def l1_loss(predictions: mx.array, targets: mx.array):
mx.array: The computed L1 loss.
"""
return mx.mean(mx.abs(predictions - targets))