docarray
diff --git a/‎docarray/computation/tensorflow_backend.py‎
Lines changed: 24 additions & 8 deletions b/‎docarray/computation/tensorflow_backend.py‎
Lines changed: 24 additions & 8 deletions
diff --git a/‎tests/units/computation_backends/tensorflow_backend/test_metrics.py‎
Lines changed: 50 additions & 0 deletions b/‎tests/units/computation_backends/tensorflow_backend/test_metrics.py‎
Lines changed: 50 additions & 0 deletions
@@ -205,10 +205,10 @@ def cosine_sim(
 
         @staticmethod
         def euclidean_dist(
-            x_mat: 'tf.Tensor',
-            y_mat: 'tf.Tensor',
+            x_mat: 'TensorFlowTensor',
+            y_mat: 'TensorFlowTensor',
             device: Optional[str] = None,
-        ) -> 'tf.Tensor':
+        ) -> 'TensorFlowTensor':
             """Pairwise Euclidian distances between all vectors in x_mat and y_mat.
 
             :param x_mat: tensor of shape (n_vectors, n_dim), where n_vectors is the
@@ -222,14 +222,25 @@ def euclidean_dist(
                 The index [i_x, i_y] contains the euclidian distance between
                 x_mat[i_x] and y_mat[i_y].
             """
-            ...
+            x_mat: tf.Tensor = TensorFlowCompBackend._norm_right(x_mat)
+            y_mat: tf.Tensor = TensorFlowCompBackend._norm_right(y_mat)
+
+            with tf.device(device):
+                x_mat, y_mat = _unsqueeze_if_single_axis(x_mat, y_mat)
+
+                dists = tf.squeeze(
+                    tf.norm(tf.subtract(x_mat, y_mat), axis=-1, ord='euclidean')
+                )
+                dists = _unsqueeze_if_scalar(dists)
+
+            return TensorFlowCompBackend._norm_left(dists)
 
         @staticmethod
         def sqeuclidean_dist(
-            x_mat: 'tf.Tensor',
-            y_mat: 'tf.Tensor',
+            x_mat: 'TensorFlowTensor',
+            y_mat: 'TensorFlowTensor',
             device: Optional[str] = None,
-        ) -> 'tf.Tensor':
+        ) -> 'TensorFlowTensor':
             """Pairwise Squared Euclidian distances between all vectors
                 in x_mat and y_mat.
 
@@ -246,4 +257,9 @@ def sqeuclidean_dist(
                 The index [i_x, i_y] contains the euclidian distance between
                 x_mat[i_x] and y_mat[i_y].
             """
-            ...
+            dists = TensorFlowCompBackend.Metrics.euclidean_dist(x_mat, y_mat)
+            squared: tf.Tensor = tf.math.square(
+                TensorFlowCompBackend._norm_right(dists)
+            )
+
+            return TensorFlowCompBackend._norm_left(squared)
@@ -19,3 +19,53 @@ def test_cosine_sim_tf():
     assert metrics.cosine_sim(b, a).tensor.shape == (5, 10)
     diag_dists = tf.linalg.diag(metrics.cosine_sim(b, b).tensor)  # self-comparisons
     tf.experimental.numpy.allclose(diag_dists, tf.ones(5))
+
+
+def test_euclidean_dist_tf():
+    a = TensorFlowTensor(tf.random.normal((128,)))
+    b = TensorFlowTensor(tf.random.normal((128,)))
+    assert metrics.euclidean_dist(a, b).tensor.shape == (1,)
+    assert metrics.euclidean_dist(a, b).tensor == metrics.euclidean_dist(b, a).tensor
+    tf.experimental.numpy.allclose(metrics.euclidean_dist(a, a).tensor, tf.zeros(1))
+
+    a = TensorFlowTensor(tf.zeros((1, 1)))
+    b = TensorFlowTensor(tf.ones((4, 1)))
+    assert metrics.euclidean_dist(a, b).tensor.shape == (4,)
+    tf.experimental.numpy.allclose(
+        metrics.euclidean_dist(a, b).tensor, metrics.euclidean_dist(b, a).tensor
+    )
+    tf.experimental.numpy.allclose(metrics.euclidean_dist(a, a).tensor, tf.zeros(1))
+
+    a = TensorFlowTensor(tf.constant([0.0, 2.0, 0.0]))
+    b = TensorFlowTensor(tf.constant([0.0, 0.0, 2.0]))
+    desired_output_singleton: tf.Tensor = tf.math.sqrt(
+        tf.constant([2.0**2.0 + 2.0**2.0])
+    )
+    tf.experimental.numpy.allclose(
+        metrics.euclidean_dist(a, b).tensor, desired_output_singleton
+    )
+
+    a = TensorFlowTensor(tf.constant([[0.0, 2.0, 0.0], [0.0, 0.0, 2.0]]))
+    b = TensorFlowTensor(tf.constant([[0.0, 0.0, 2.0], [0.0, 2.0, 0.0]]))
+    desired_output_singleton = tf.constant([[2.828427, 0.0], [0.0, 2.828427]])
+    tf.experimental.numpy.allclose(
+        metrics.euclidean_dist(a, b).tensor, desired_output_singleton
+    )
+
+
+def test_sqeuclidean_dist_torch():
+    a = TensorFlowTensor(tf.random.normal((128,)))
+    b = TensorFlowTensor(tf.random.normal((128,)))
+    assert metrics.sqeuclidean_dist(a, b).tensor.shape == (1,)
+    tf.experimental.numpy.allclose(
+        metrics.sqeuclidean_dist(a, b).tensor,
+        metrics.euclidean_dist(a, b).tensor ** 2,
+    )
+
+    a = TensorFlowTensor(tf.random.normal((1, 1)))
+    b = TensorFlowTensor(tf.random.normal((4, 1)))
+    assert metrics.sqeuclidean_dist(b, a).tensor.shape == (4,)
+    tf.experimental.numpy.allclose(
+        metrics.sqeuclidean_dist(a, b).tensor,
+        metrics.euclidean_dist(a, b).tensor ** 2,
+    )