zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-21 12:06:42 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Update test_nn.py

Browse Source 
cos sim test
This commit is contained in:
YUN, Junwoo 2024-01-02 23:11:46 +09:00 committed by GitHub
parent d66fec6e81
commit d408541869
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 28 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
python/tests/test_nn.py
View File

@ -265,6 +265,32 @@ class TestNN(mlx_tests.MLXTestCase):
expected_sum = mx.sum(expected_none)
self.assertTrue(mx.allclose(losses_sum, expected_sum))
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]])
targets = mx.array([1, -1])
# Test with reduction 'none'
losses_none = nn.losses.cosine_similarity_loss(
embeddings1, embeddings2, targets, reduction="none"
)
expected_none = mx.array([0.0146555, 0.961074])
self.assertTrue(mx.allclose(losses_none, expected_none))
# Test with reduction 'mean'
losses_mean = nn.losses.cosine_similarity_loss(
embeddings1, embeddings2, targets, 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, targets, reduction="sum"
)
expected_sum = mx.sum(expected_none)
self.assertTrue(mx.allclose(losses_sum, expected_sum))
def test_gelu(self):
inputs = [1.15286231, -0.81037411, 0.35816911, 0.77484438, 0.66276414]
@ -843,3 +869,4 @@ class TestNN(mlx_tests.MLXTestCase):
if __name__ == "__main__":
unittest.main()