mirror of
https://github.com/ml-explore/mlx.git
synced 2025-10-19 08:38:09 +08:00
Additoinal losses (#336)
* cosine similarity loss --------- Co-authored-by: Awni Hannun <awni@apple.com> * Docstring nits
This commit is contained in:
@@ -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()
|
||||
|
Reference in New Issue
Block a user