From 644e4d9f26da6c06167bcbac374e293a57439e6c Mon Sep 17 00:00:00 2001 From: Aziz Belaweid Date: Mon, 21 Mar 2022 15:24:36 +0100 Subject: [PATCH 1/3] test(distance): unit tests for distance --- tests/unit/math/distance/test_numpy.py | 121 ++++++++++++++++++++ tests/unit/math/distance/test_paddle.py | 94 +++++++++++++++ tests/unit/math/distance/test_tensorflow.py | 94 +++++++++++++++ tests/unit/math/distance/test_torch.py | 94 +++++++++++++++ 4 files changed, 403 insertions(+) create mode 100644 tests/unit/math/distance/test_numpy.py create mode 100644 tests/unit/math/distance/test_paddle.py create mode 100644 tests/unit/math/distance/test_tensorflow.py create mode 100644 tests/unit/math/distance/test_torch.py diff --git a/tests/unit/math/distance/test_numpy.py b/tests/unit/math/distance/test_numpy.py new file mode 100644 index 00000000000..b48086da420 --- /dev/null +++ b/tests/unit/math/distance/test_numpy.py @@ -0,0 +1,121 @@ +import numpy as np +import pytest +from scipy.sparse import csr_matrix + +from docarray.math.distance.numpy import (cosine, euclidean, sparse_cosine, + sparse_euclidean, sparse_sqeuclidean, + sqeuclidean) + + +@pytest.mark.parametrize( + 'x_mat, y_mat, result', + ( + ( + np.array([[1, 2, 3], [4, 5, 6]]), + np.array([[1, 2, 3], [4, 5, 6]]), + np.array( + [[0.00000000e00, 2.53681537e-02], [2.53681537e-02, 2.22044605e-16]] + ), + ), + (np.array([[1, 2, 3]]), np.array([[1, 2, 3]]), np.array([[0]])), + (np.array([[0, 0, 0]]), np.array([[0, 0, 0]]), np.array([[0]])), + (np.array([[1, 2, 3]]), np.array([[19, 53, 201]]), np.array([[0.06788693]])), + ), +) +def test_cosine(x_mat, y_mat, result): + assert cosine(x_mat, y_mat).all() == result.all() + + +@pytest.mark.parametrize( + 'x_mat, y_mat, result', + ( + ( + csr_matrix([[1, 2, 3], [4, 5, 6]]), + csr_matrix([[1, 2, 3], [4, 5, 6]]), + np.array( + [[0.00000000e00, 2.53681537e-02], [2.53681537e-02, 2.22044605e-16]] + ), + ), + (csr_matrix([[1, 2, 3]]), csr_matrix([[1, 2, 3]]), np.array([[0]])), + (csr_matrix([[0, 0, 0]]), csr_matrix([[0, 0, 0]]), np.array([[np.nan]])), + ( + csr_matrix([[1, 2, 3]]), + csr_matrix([[19, 53, 201]]), + np.array([[0.06788693]]), + ), + ), +) +def test_sparse_cosine(x_mat, y_mat, result): + assert sparse_cosine(x_mat, y_mat).all() == result.all() + + +@pytest.mark.parametrize( + 'x_mat, y_mat, result', + ( + ( + np.array([[1, 2, 3], [4, 5, 6]]), + np.array([[1, 2, 3], [4, 5, 6]]), + np.array([[0, 27], [27, 0]]), + ), + (np.array([[1, 2, 3]]), np.array([[1, 2, 3]]), np.array([[0]])), + (np.array([[0, 0, 0]]), np.array([[0, 0, 0]]), np.array([[0]])), + (np.array([[1, 2, 3]]), np.array([[19, 53, 201]]), np.array([[42129]])), + ), +) +def test_sqeuclidean(x_mat, y_mat, result): + assert sqeuclidean(x_mat, y_mat).all() == result.all() + + +@pytest.mark.parametrize( + 'x_mat, y_mat, result', + ( + ( + csr_matrix([[1, 2, 3], [4, 5, 6]]), + csr_matrix([[1, 2, 3], [4, 5, 6]]), + np.array([[0, 27], [27, 0]]), + ), + (csr_matrix([[1, 2, 3]]), csr_matrix([[1, 2, 3]]), np.array([[0]])), + (csr_matrix([[0, 0, 0]]), csr_matrix([[0, 0, 0]]), np.array([[0]])), + (csr_matrix([[1, 2, 3]]), csr_matrix([[19, 53, 201]]), np.array([[42129]])), + ), +) +def test_sparse_sqeuclidean(x_mat, y_mat, result): + assert sparse_sqeuclidean(x_mat, y_mat).all() == result.all() + + +@pytest.mark.parametrize( + 'x_mat, y_mat, result', + ( + ( + np.array([[1, 2, 3], [4, 5, 6]]), + np.array([[1, 2, 3], [4, 5, 6]]), + np.array([[0, 5.19615242], [5.19615242, 0]]), + ), + (np.array([[1, 2, 3]]), np.array([[1, 2, 3]]), np.array([[0]])), + (np.array([[0, 0, 0]]), np.array([[0, 0, 0]]), np.array([[0]])), + (np.array([[1, 2, 3]]), np.array([[19, 53, 201]]), np.array([[205.2535018]])), + ), +) +def test_euclidean(x_mat, y_mat, result): + assert euclidean(x_mat, y_mat).all() == result.all() + + +@pytest.mark.parametrize( + 'x_mat, y_mat, result', + ( + ( + csr_matrix([[1, 2, 3], [4, 5, 6]]), + csr_matrix([[1, 2, 3], [4, 5, 6]]), + np.array([[0, 5.19615242], [5.19615242, 0]]), + ), + (csr_matrix([[1, 2, 3]]), csr_matrix([[1, 2, 3]]), np.array([[0]])), + (csr_matrix([[0, 0, 0]]), csr_matrix([[0, 0, 0]]), np.array([[0]])), + ( + csr_matrix([[1, 2, 3]]), + csr_matrix([[19, 53, 201]]), + np.array([[205.2535018]]), + ), + ), +) +def test_sparse_euclidean(x_mat, y_mat, result): + assert sparse_euclidean(x_mat, y_mat).all() == result.all() diff --git a/tests/unit/math/distance/test_paddle.py b/tests/unit/math/distance/test_paddle.py new file mode 100644 index 00000000000..e7b0e502295 --- /dev/null +++ b/tests/unit/math/distance/test_paddle.py @@ -0,0 +1,94 @@ +import numpy as np +import paddle +import pytest + +from docarray.math.distance.paddle import cosine, euclidean, sqeuclidean + + +@pytest.mark.parametrize( + 'x_mat, y_mat, result', + ( + ( + paddle.to_tensor([[1, 2, 3], [4, 5, 6]], dtype='float32'), + paddle.to_tensor([[1, 2, 3], [4, 5, 6]], dtype='float32'), + np.array( + [[0.00000000e00, 2.53681537e-02], [2.53681537e-02, 2.22044605e-16]] + ), + ), + ( + paddle.to_tensor([[1, 2, 3]], dtype='float32'), + paddle.to_tensor([[1, 2, 3]], dtype='float32'), + np.array([[1]]), + ), + ( + paddle.to_tensor([[0, 0, 0]], dtype='float32'), + paddle.to_tensor([[0, 0, 0]], dtype='float32'), + np.array([[1]]), + ), + ( + paddle.to_tensor([[1, 2, 3]], dtype='float32'), + paddle.to_tensor([[19, 53, 201]], dtype='float32'), + np.array([[0.06788693]]), + ), + ), +) +def test_cosine(x_mat, y_mat, result): + assert cosine(x_mat, y_mat).all() == result.all() + + +@pytest.mark.parametrize( + 'x_mat, y_mat, result', + ( + ( + paddle.to_tensor([[1, 2, 3], [4, 5, 6]], dtype='float32'), + paddle.to_tensor([[1, 2, 3], [4, 5, 6]], dtype='float32'), + np.array([[0, 27], [27, 0]]), + ), + ( + paddle.to_tensor([[1, 2, 3]], dtype='float32'), + paddle.to_tensor([[1, 2, 3]], dtype='float32'), + np.array([[0]]), + ), + ( + paddle.to_tensor([[0, 0, 0]], dtype='float32'), + paddle.to_tensor([[0, 0, 0]], dtype='float32'), + np.array([[0]]), + ), + ( + paddle.to_tensor([[1, 2, 3]], dtype='float32'), + paddle.to_tensor([[19, 53, 201]], dtype='float32'), + np.array([[42129]]), + ), + ), +) +def test_sqeuclidean(x_mat, y_mat, result): + assert sqeuclidean(x_mat, y_mat).all() == result.all() + + +@pytest.mark.parametrize( + 'x_mat, y_mat, result', + ( + ( + paddle.to_tensor([[1, 2, 3], [4, 5, 6]], dtype='float32'), + paddle.to_tensor([[1, 2, 3], [4, 5, 6]], dtype='float32'), + np.array([[0, 5.19615242], [5.19615242, 0]]), + ), + ( + paddle.to_tensor([[1, 2, 3]], dtype='float32'), + paddle.to_tensor([[1, 2, 3]], dtype='float32'), + np.array([[0]]), + ), + ( + paddle.to_tensor([[0, 0, 0]], dtype='float32'), + paddle.to_tensor([[0, 0, 0]], dtype='float32'), + np.array([[0]]), + ), + ( + paddle.to_tensor([[1, 2, 3]], dtype='float32'), + paddle.to_tensor([[19, 53, 201]], dtype='float32'), + np.array([[205.2535018]]), + ), + ), +) +def test_euclidean(x_mat, y_mat, result): + assert euclidean(x_mat, y_mat).all() == result.all() diff --git a/tests/unit/math/distance/test_tensorflow.py b/tests/unit/math/distance/test_tensorflow.py new file mode 100644 index 00000000000..dff4c261133 --- /dev/null +++ b/tests/unit/math/distance/test_tensorflow.py @@ -0,0 +1,94 @@ +import numpy as np +import pytest +import tensorflow as tf + +from docarray.math.distance.tensorflow import cosine, euclidean, sqeuclidean + + +@pytest.mark.parametrize( + 'x_mat, y_mat, result', + ( + ( + tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.float32), + tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.float32), + np.array( + [[0.00000000e00, 2.53681537e-02], [2.53681537e-02, 2.22044605e-16]] + ), + ), + ( + tf.constant([[1, 2, 3]], dtype=tf.float32), + tf.constant([[1, 2, 3]], dtype=tf.float32), + np.array([[1]]), + ), + ( + tf.constant([[0, 0, 0]], dtype=tf.float32), + tf.constant([[0, 0, 0]], dtype=tf.float32), + np.array([[1]]), + ), + ( + tf.constant([[1, 2, 3]], dtype=tf.float32), + tf.constant([[19, 53, 201]], dtype=tf.float32), + np.array([[0.06788693]]), + ), + ), +) +def test_cosine(x_mat, y_mat, result): + assert cosine(x_mat, y_mat).all() == result.all() + + +@pytest.mark.parametrize( + 'x_mat, y_mat, result', + ( + ( + tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.float32), + tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.float32), + np.array([[0, 27], [27, 0]]), + ), + ( + tf.constant([[1, 2, 3]], dtype=tf.float32), + tf.constant([[1, 2, 3]], dtype=tf.float32), + np.array([[0]]), + ), + ( + tf.constant([[0, 0, 0]], dtype=tf.float32), + tf.constant([[0, 0, 0]], dtype=tf.float32), + np.array([[0]]), + ), + ( + tf.constant([[1, 2, 3]], dtype=tf.float32), + tf.constant([[19, 53, 201]], dtype=tf.float32), + np.array([[42129]]), + ), + ), +) +def test_sqeuclidean(x_mat, y_mat, result): + assert sqeuclidean(x_mat, y_mat).all() == result.all() + + +@pytest.mark.parametrize( + 'x_mat, y_mat, result', + ( + ( + tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.float32), + tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.float32), + np.array([[0, 5.19615242], [5.19615242, 0]]), + ), + ( + tf.constant([[1, 2, 3]], dtype=tf.float32), + tf.constant([[1, 2, 3]], dtype=tf.float32), + np.array([[0]]), + ), + ( + tf.constant([[0, 0, 0]], dtype=tf.float32), + tf.constant([[0, 0, 0]], dtype=tf.float32), + np.array([[0]]), + ), + ( + tf.constant([[1, 2, 3]], dtype=tf.float32), + tf.constant([[19, 53, 201]], dtype=tf.float32), + np.array([[205.2535018]]), + ), + ), +) +def test_euclidean(x_mat, y_mat, result): + assert euclidean(x_mat, y_mat).all() == result.all() diff --git a/tests/unit/math/distance/test_torch.py b/tests/unit/math/distance/test_torch.py new file mode 100644 index 00000000000..35f8541f6d8 --- /dev/null +++ b/tests/unit/math/distance/test_torch.py @@ -0,0 +1,94 @@ +import numpy as np +import pytest +import torch + +from docarray.math.distance.torch import cosine, euclidean, sqeuclidean + + +@pytest.mark.parametrize( + 'x_mat, y_mat, result', + ( + ( + torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), + torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), + np.array( + [[0.00000000e00, 2.53681537e-02], [2.53681537e-02, 2.22044605e-16]] + ), + ), + ( + torch.tensor([[1.0, 2.0, 3.0]]), + torch.tensor([[1.0, 2.0, 3.0]]), + np.array([[1]]), + ), + ( + torch.tensor([[0.0, 0.0, 0.0]]), + torch.tensor([[0.0, 0.0, 0.0]]), + np.array([[1]]), + ), + ( + torch.tensor([[1.0, 2.0, 3.0]]), + torch.tensor([[19.0, 53.0, 201.0]]), + np.array([[0.06788693]]), + ), + ), +) +def test_cosine(x_mat, y_mat, result): + assert cosine(x_mat, y_mat).all() == result.all() + + +@pytest.mark.parametrize( + 'x_mat, y_mat, result', + ( + ( + torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), + torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), + np.array([[0, 27], [27, 0]]), + ), + ( + torch.tensor([[1.0, 2.0, 3.0]]), + torch.tensor([[1.0, 2.0, 3.0]]), + np.array([[0]]), + ), + ( + torch.tensor([[0.0, 0.0, 0.0]]), + torch.tensor([[0.0, 0.0, 0.0]]), + np.array([[0]]), + ), + ( + torch.tensor([[1.0, 2.0, 3.0]]), + torch.tensor([[19.0, 53.0, 201.0]]), + np.array([[42129]]), + ), + ), +) +def test_sqeuclidean(x_mat, y_mat, result): + assert sqeuclidean(x_mat, y_mat).all() == result.all() + + +@pytest.mark.parametrize( + 'x_mat, y_mat, result', + ( + ( + torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), + torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), + np.array([[0, 5.19615242], [5.19615242, 0]]), + ), + ( + torch.tensor([[1.0, 2.0, 3.0]]), + torch.tensor([[1.0, 2.0, 3.0]]), + np.array([[0]]), + ), + ( + torch.tensor([[0.0, 0.0, 0.0]]), + torch.tensor([[0.0, 0.0, 0.0]]), + np.array([[0]]), + ), + ( + torch.tensor([[1.0, 2.0, 3.0]]), + torch.tensor([[19.0, 53.0, 201.0]]), + np.array([[205.2535018]]), + ), + ), +) +def test_euclidean(x_mat, y_mat, result): + assert euclidean(x_mat, y_mat).all() == result.all() From 3991f678adc3aa8dc238d7e884b85dcb739a6fe7 Mon Sep 17 00:00:00 2001 From: Aziz Belaweid Date: Mon, 21 Mar 2022 15:27:17 +0100 Subject: [PATCH 2/3] style: black --- tests/unit/math/distance/test_numpy.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/unit/math/distance/test_numpy.py b/tests/unit/math/distance/test_numpy.py index b48086da420..d964194bb90 100644 --- a/tests/unit/math/distance/test_numpy.py +++ b/tests/unit/math/distance/test_numpy.py @@ -2,9 +2,14 @@ import pytest from scipy.sparse import csr_matrix -from docarray.math.distance.numpy import (cosine, euclidean, sparse_cosine, - sparse_euclidean, sparse_sqeuclidean, - sqeuclidean) +from docarray.math.distance.numpy import ( + cosine, + euclidean, + sparse_cosine, + sparse_euclidean, + sparse_sqeuclidean, + sqeuclidean, +) @pytest.mark.parametrize( From ab6252d4c7c525ec67956b01f83d5967dee6b906 Mon Sep 17 00:00:00 2001 From: Aziz Belaweid Date: Mon, 21 Mar 2022 16:21:12 +0100 Subject: [PATCH 3/3] feat: use np testing --- tests/unit/math/distance/test_numpy.py | 12 ++++++------ tests/unit/math/distance/test_paddle.py | 12 +++++------- tests/unit/math/distance/test_tensorflow.py | 12 +++++------- tests/unit/math/distance/test_torch.py | 12 +++++------- 4 files changed, 21 insertions(+), 27 deletions(-) diff --git a/tests/unit/math/distance/test_numpy.py b/tests/unit/math/distance/test_numpy.py index d964194bb90..473893a296d 100644 --- a/tests/unit/math/distance/test_numpy.py +++ b/tests/unit/math/distance/test_numpy.py @@ -28,7 +28,7 @@ ), ) def test_cosine(x_mat, y_mat, result): - assert cosine(x_mat, y_mat).all() == result.all() + np.testing.assert_allclose(cosine(x_mat, y_mat), result, rtol=1e-5) @pytest.mark.parametrize( @@ -51,7 +51,7 @@ def test_cosine(x_mat, y_mat, result): ), ) def test_sparse_cosine(x_mat, y_mat, result): - assert sparse_cosine(x_mat, y_mat).all() == result.all() + np.testing.assert_allclose(sparse_cosine(x_mat, y_mat), result, rtol=1e-5) @pytest.mark.parametrize( @@ -68,7 +68,7 @@ def test_sparse_cosine(x_mat, y_mat, result): ), ) def test_sqeuclidean(x_mat, y_mat, result): - assert sqeuclidean(x_mat, y_mat).all() == result.all() + np.testing.assert_allclose(sqeuclidean(x_mat, y_mat), result, rtol=1e-5) @pytest.mark.parametrize( @@ -85,7 +85,7 @@ def test_sqeuclidean(x_mat, y_mat, result): ), ) def test_sparse_sqeuclidean(x_mat, y_mat, result): - assert sparse_sqeuclidean(x_mat, y_mat).all() == result.all() + np.testing.assert_allclose(sparse_sqeuclidean(x_mat, y_mat), result, rtol=1e-5) @pytest.mark.parametrize( @@ -102,7 +102,7 @@ def test_sparse_sqeuclidean(x_mat, y_mat, result): ), ) def test_euclidean(x_mat, y_mat, result): - assert euclidean(x_mat, y_mat).all() == result.all() + np.testing.assert_allclose(euclidean(x_mat, y_mat), result, rtol=1e-5) @pytest.mark.parametrize( @@ -123,4 +123,4 @@ def test_euclidean(x_mat, y_mat, result): ), ) def test_sparse_euclidean(x_mat, y_mat, result): - assert sparse_euclidean(x_mat, y_mat).all() == result.all() + np.testing.assert_allclose(sparse_euclidean(x_mat, y_mat), result, rtol=1e-5) diff --git a/tests/unit/math/distance/test_paddle.py b/tests/unit/math/distance/test_paddle.py index e7b0e502295..2f33b182a16 100644 --- a/tests/unit/math/distance/test_paddle.py +++ b/tests/unit/math/distance/test_paddle.py @@ -11,14 +11,12 @@ ( paddle.to_tensor([[1, 2, 3], [4, 5, 6]], dtype='float32'), paddle.to_tensor([[1, 2, 3], [4, 5, 6]], dtype='float32'), - np.array( - [[0.00000000e00, 2.53681537e-02], [2.53681537e-02, 2.22044605e-16]] - ), + np.array([[1.192093e-07, 2.53681537e-02], [2.53681537e-02, 0]]), ), ( paddle.to_tensor([[1, 2, 3]], dtype='float32'), paddle.to_tensor([[1, 2, 3]], dtype='float32'), - np.array([[1]]), + np.array([[1.192093e-07]]), ), ( paddle.to_tensor([[0, 0, 0]], dtype='float32'), @@ -33,7 +31,7 @@ ), ) def test_cosine(x_mat, y_mat, result): - assert cosine(x_mat, y_mat).all() == result.all() + np.testing.assert_allclose(cosine(x_mat, y_mat), result, rtol=1e-5) @pytest.mark.parametrize( @@ -62,7 +60,7 @@ def test_cosine(x_mat, y_mat, result): ), ) def test_sqeuclidean(x_mat, y_mat, result): - assert sqeuclidean(x_mat, y_mat).all() == result.all() + np.testing.assert_allclose(sqeuclidean(x_mat, y_mat), result, rtol=1e-5) @pytest.mark.parametrize( @@ -91,4 +89,4 @@ def test_sqeuclidean(x_mat, y_mat, result): ), ) def test_euclidean(x_mat, y_mat, result): - assert euclidean(x_mat, y_mat).all() == result.all() + np.testing.assert_allclose(euclidean(x_mat, y_mat), result, rtol=1e-5) diff --git a/tests/unit/math/distance/test_tensorflow.py b/tests/unit/math/distance/test_tensorflow.py index dff4c261133..4578005479a 100644 --- a/tests/unit/math/distance/test_tensorflow.py +++ b/tests/unit/math/distance/test_tensorflow.py @@ -11,14 +11,12 @@ ( tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.float32), tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.float32), - np.array( - [[0.00000000e00, 2.53681537e-02], [2.53681537e-02, 2.22044605e-16]] - ), + np.array([[1.192093e-07, 2.53681537e-02], [2.53681537e-02, 0.000000e00]]), ), ( tf.constant([[1, 2, 3]], dtype=tf.float32), tf.constant([[1, 2, 3]], dtype=tf.float32), - np.array([[1]]), + np.array([[1.192093e-07]]), ), ( tf.constant([[0, 0, 0]], dtype=tf.float32), @@ -33,7 +31,7 @@ ), ) def test_cosine(x_mat, y_mat, result): - assert cosine(x_mat, y_mat).all() == result.all() + np.testing.assert_allclose(cosine(x_mat, y_mat), result, rtol=1e-5) @pytest.mark.parametrize( @@ -62,7 +60,7 @@ def test_cosine(x_mat, y_mat, result): ), ) def test_sqeuclidean(x_mat, y_mat, result): - assert sqeuclidean(x_mat, y_mat).all() == result.all() + np.testing.assert_allclose(sqeuclidean(x_mat, y_mat), result, rtol=1e-5) @pytest.mark.parametrize( @@ -91,4 +89,4 @@ def test_sqeuclidean(x_mat, y_mat, result): ), ) def test_euclidean(x_mat, y_mat, result): - assert euclidean(x_mat, y_mat).all() == result.all() + np.testing.assert_allclose(euclidean(x_mat, y_mat), result, rtol=1e-5) diff --git a/tests/unit/math/distance/test_torch.py b/tests/unit/math/distance/test_torch.py index 35f8541f6d8..f769d2402ed 100644 --- a/tests/unit/math/distance/test_torch.py +++ b/tests/unit/math/distance/test_torch.py @@ -11,14 +11,12 @@ ( torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), - np.array( - [[0.00000000e00, 2.53681537e-02], [2.53681537e-02, 2.22044605e-16]] - ), + np.array([[1.192093e-07, 2.53681537e-02], [2.53681537e-02, 0.000000e00]]), ), ( torch.tensor([[1.0, 2.0, 3.0]]), torch.tensor([[1.0, 2.0, 3.0]]), - np.array([[1]]), + np.array([[1.192093e-07]], dtype=np.float32), ), ( torch.tensor([[0.0, 0.0, 0.0]]), @@ -33,7 +31,7 @@ ), ) def test_cosine(x_mat, y_mat, result): - assert cosine(x_mat, y_mat).all() == result.all() + np.testing.assert_allclose(cosine(x_mat, y_mat), result, rtol=1e-5) @pytest.mark.parametrize( @@ -62,7 +60,7 @@ def test_cosine(x_mat, y_mat, result): ), ) def test_sqeuclidean(x_mat, y_mat, result): - assert sqeuclidean(x_mat, y_mat).all() == result.all() + np.testing.assert_allclose(sqeuclidean(x_mat, y_mat), result, rtol=1e-5) @pytest.mark.parametrize( @@ -91,4 +89,4 @@ def test_sqeuclidean(x_mat, y_mat, result): ), ) def test_euclidean(x_mat, y_mat, result): - assert euclidean(x_mat, y_mat).all() == result.all() + np.testing.assert_allclose(euclidean(x_mat, y_mat), result, rtol=1e-5)