Introduction

Pairwise distance means building an N x M matrix in which each entry compares one vector from set A with one vector from set B. TensorFlow can do this efficiently, but the main design choice is whether you want the simplest broadcasting code or the lower-memory matrix formula.

The efficient formula for squared Euclidean distance

For Euclidean distance, a common trick avoids explicitly materializing every pairwise difference vector. Use the identity based on vector norms and a matrix multiply.

1import tensorflow as tf
2
3def pairwise_squared_distance(a: tf.Tensor, b: tf.Tensor) -> tf.Tensor:
4    a_sq = tf.reduce_sum(tf.square(a), axis=1, keepdims=True)
5    b_sq = tf.reduce_sum(tf.square(b), axis=1, keepdims=True)
6
7    distances = a_sq - 2.0 * tf.matmul(a, b, transpose_b=True) + tf.transpose(b_sq)
8    return tf.maximum(distances, 0.0)
9
10
11a = tf.constant([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]], dtype=tf.float32)
12b = tf.constant([[1.0, 0.0], [2.0, 1.0]], dtype=tf.float32)
13
14print(pairwise_squared_distance(a, b).numpy())

This returns squared distances, which are often enough for nearest-neighbor ranking or loss functions. The tf.maximum call guards against tiny negative values caused by floating-point roundoff.

Converting to true Euclidean distance

If you really need Euclidean distance rather than squared distance, apply a square root at the end.

distances = tf.sqrt(pairwise_squared_distance(a, b))
print(distances.numpy())

Delaying the square root is useful because many algorithms only need relative ordering. Skipping sqrt saves work and avoids one more source of numerical noise.

Broadcasting is simpler but heavier

TensorFlow broadcasting can express pairwise distances more directly:

def pairwise_distance_broadcast(a: tf.Tensor, b: tf.Tensor) -> tf.Tensor:
    diff = a[:, None, :] - b[None, :, :]
    return tf.sqrt(tf.reduce_sum(tf.square(diff), axis=-1))

This version is easy to read, but it materializes an intermediate tensor shaped like N x M x D. That can become expensive quickly when the number of points or feature dimensions grows. The matrix-multiply formulation is usually the better default for larger inputs.

Choosing the right TensorFlow pattern

If you are comparing a moderate number of vectors and want clarity, broadcasting is fine. If you are working with many vectors, use the squared-distance formula first. If N x M itself is too large to fit comfortably in memory, batch the computation by splitting one input set into chunks and concatenating the results later.

TensorFlow gives you GPU acceleration automatically once the tensors live on a GPU-capable device, but acceleration does not change the fact that a huge pairwise matrix still consumes huge memory.

If you are comparing one set against itself, remember that the diagonal should be zero in exact arithmetic. Small nonzero values on the diagonal are usually just floating-point residue, which is another reason to clamp near-zero negatives before taking a square root or applying downstream logic that expects exact self-distance behavior.

Common Pitfalls

• Building the full broadcasted difference tensor for very large inputs and running out of memory.
• Taking square roots when squared distance would have been sufficient.
• Forgetting to clamp tiny negative values caused by floating-point error before applying sqrt.
• Mixing tensors with different dtypes and getting unexpected casts or slower execution.
• Assuming that GPU acceleration removes the need to think about N x M output size.

Summary

• Use the norm-plus-matmul identity for efficient pairwise squared Euclidean distance.
• Apply sqrt only if you need true Euclidean distance values.
• Broadcasting is readable, but it uses more memory because it materializes pairwise differences.
• Clamp small negative values to zero before square-rooting.
• The output matrix size still dominates memory use, even when TensorFlow runs the math quickly.