test: test when device is cuda (docarray#293)

azayz · alaeddine-13 · web-flow · commit ea98df641fb0 · 2022-04-25T10:04:16.000+02:00
Co-authored-by: AlaeddineAbdessalem &lt;alaeddine-13@live.fr&gt;
diff --git a/tests/unit/math/distance/test_distance.py b/tests/unit/math/distance/test_distance.py
@@ -0,0 +1,82 @@
+import numpy as np
+import paddle
+import pytest
+import tensorflow as tf
+import torch
+
+from docarray.math.distance import cdist, pdist
+
+
+def test_pdist():
+    tensor = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+    np.testing.assert_almost_equal(
+        pdist(tensor, 'cosine'),
+        cdist(tensor, tensor, 'cosine'),
+        decimal=3,
+    )
+
+
+def test_cdist_raise_error():
+    x = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+    y = np.array([[1, 2, 3], [4, 5, 6]])
+    with pytest.raises(ValueError):
+        cdist(x, y, 'cosine')
+
+
+def test_not_supported_metric():
+    x = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+    y = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+    with pytest.raises(NotImplementedError):
+        cdist(x, y, 'fake_metric')
+
+
+@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]]),
+        ),
+        (
+            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]]),
+        ),
+        (
+            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]]),
+        ),
+    ),
+)
+def test_seqeuclidean(x_mat, y_mat, result):
+    np.testing.assert_almost_equal(
+        cdist(x_mat, y_mat, metric='sqeuclidean'), result, decimal=3
+    )
+
+
+@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.196], [5.196, 0]]),
+        ),
+        (
+            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.196], [5.196, 0]]),
+        ),
+        (
+            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.196], [5.196, 0]]),
+        ),
+    ),
+)
+def test_euclidean(x_mat, y_mat, result):
+    np.testing.assert_almost_equal(
+        cdist(x_mat, y_mat, metric='euclidean'), result, decimal=3
+    )
