zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-03 22:34:43 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

MLE and L1 loss functions (#88)

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
This commit is contained in:
Kai Ma
2023-12-08 23:21:37 -05:00
committed by GitHub
parent 2b714714e1
commit 641d316484
2 changed files with 61 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
python/tests/test_nn.py
View File

@@ -10,7 +10,6 @@ import mlx_tests
import numpy as np
from mlx.utils import tree_flatten, tree_map, tree_unflatten
class TestNN(mlx_tests.MLXTestCase):
def test_linear(self):
inputs = mx.zeros((10, 4))
@@ -21,8 +20,27 @@ class TestNN(mlx_tests.MLXTestCase):
def test_cross_entropy(self):
logits = mx.array([[0.0, -float("inf")], [-float("inf"), 0.0]])
targets = mx.array([0, 1])
losses = nn.losses.cross_entropy(logits, targets)
self.assertTrue(mx.array_equal(losses, mx.zeros((2,))))
# Test with reduction 'none'
losses_none = nn.losses.cross_entropy(logits, targets, reduction='none')
expected_none = mx.array([0.0, 0.0])
self.assertTrue(mx.array_equal(losses_none, expected_none))
# Test with reduction 'mean'
losses_mean = nn.losses.cross_entropy(logits, targets, reduction='mean')
expected_mean = mx.mean(expected_none)
self.assertEqual(losses_mean, expected_mean)
# Test with reduction 'sum'
losses_sum = nn.losses.cross_entropy(logits, targets, reduction='sum')
expected_sum = mx.sum(expected_none)
self.assertEqual(losses_sum, expected_sum)
def test_l1_loss(self):
predictions = mx.array([0.5, 0.2, 0.9, 0.0])
targets = mx.array([0.5, 0.2, 0.9, 0.0])
losses = nn.losses.l1_loss(predictions, targets)
self.assertEqual(losses, 0.0)
def test_gelu(self):
inputs = [1.15286231, -0.81037411, 0.35816911, 0.77484438, 0.66276414]