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

Additoinal losses (#336)

Browse Source 
* cosine similarity loss

---------

Co-authored-by: Awni Hannun <awni@apple.com>

* Docstring nits
This commit is contained in:
YUN, Junwoo
2024-01-09 07:01:13 +09:00
committed by GitHub
parent 432ee5650b
commit 0b8aeddac6
3 changed files with 74 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
python/tests/test_losses.py
View File

@@ -274,6 +274,31 @@ class TestLosses(mlx_tests.MLXTestCase):
loss = nn.losses.log_cosh_loss(inputs, targets, reduction="mean")
self.assertAlmostEqual(loss.item(), 0.433781, places=6)
def test_cosine_similarity_loss(self):
embeddings1 = mx.array([[0.5, 0.5, 0.2, 0.9], [0.1, 0.3, 0.5, 0.5]])
embeddings2 = mx.array([[0.6, 0.4, 0.3, 0.8], [0.2, 0.5, 0.6, 0.4]])
# Test with reduction 'none'
losses_none = nn.losses.cosine_similarity_loss(
embeddings1, embeddings2, reduction="none"
)
expected_none = mx.array([0.985344, 0.961074])
self.assertTrue(mx.allclose(losses_none, expected_none))
# Test with reduction 'mean'
losses_mean = nn.losses.cosine_similarity_loss(
embeddings1, embeddings2, reduction="mean"
)
expected_mean = mx.mean(expected_none)
self.assertTrue(mx.allclose(losses_mean, expected_mean))
# Test with reduction 'sum'
losses_sum = nn.losses.cosine_similarity_loss(
embeddings1, embeddings2, reduction="sum"
)
expected_sum = mx.sum(expected_none)
self.assertTrue(mx.allclose(losses_sum, expected_sum))
if __name__ == "__main__":
unittest.main()