Introduction
To evaluate all pairwise combinations of rows between two tensors in TensorFlow, use broadcasting by expanding dimensions with tf.expand_dims so that an (M, D) tensor and an (N, D) tensor produce an (M, N, ...) result. This avoids explicit loops and leverages GPU-accelerated batch operations. Common applications include computing pairwise distances, dot products, or applying a learned similarity function between all pairs of data points.
The Broadcasting Approach
Given tensor A with shape (M, D) and tensor B with shape (N, D), expand their dimensions to enable broadcasting:
1import tensorflow as tf
2
3A = tf.constant([[1.0, 2.0],
4 [3.0, 4.0],
5 [5.0, 6.0]]) # Shape: (3, 2)
6
7B = tf.constant([[10.0, 20.0],
8 [30.0, 40.0]]) # Shape: (2, 2)
9
10# Expand dimensions for broadcasting
11A_expanded = tf.expand_dims(A, axis=1) # Shape: (3, 1, 2)
12B_expanded = tf.expand_dims(B, axis=0) # Shape: (1, 2, 2)
13
14# Now any element-wise operation broadcasts to (3, 2, 2)
15pairwise_sum = A_expanded + B_expanded
16print(pairwise_sum.shape) # (3, 2, 2)
17# pairwise_sum[i, j] = A[i] + B[j]
Pairwise Euclidean Distance
1def pairwise_euclidean_distance(A, B):
2 """Compute distance between every row of A and every row of B.
3
4 Args:
5 A: Tensor of shape (M, D)
6 B: Tensor of shape (N, D)
7
8 Returns:
9 Tensor of shape (M, N) where result[i, j] = ||A[i] - B[j]||
10 """
11 A_expanded = tf.expand_dims(A, axis=1) # (M, 1, D)
12 B_expanded = tf.expand_dims(B, axis=0) # (1, N, D)
13
14 diff = A_expanded - B_expanded # (M, N, D)
15 distances = tf.sqrt(tf.reduce_sum(tf.square(diff), axis=-1)) # (M, N)
16 return distances
17
18A = tf.constant([[0.0, 0.0], [3.0, 4.0]])
19B = tf.constant([[1.0, 0.0], [0.0, 1.0], [3.0, 4.0]])
20
21dist = pairwise_euclidean_distance(A, B)
22print(dist)
23# [[1.0, 1.0, 5.0],
24# [3.6, 4.2, 0.0]]
Squared Distance (More Efficient)
1def pairwise_squared_distance(A, B):
2 """Compute squared Euclidean distance without sqrt."""
3 # ||a - b||^2 = ||a||^2 - 2*a·b + ||b||^2
4 A_sq = tf.reduce_sum(tf.square(A), axis=-1, keepdims=True) # (M, 1)
5 B_sq = tf.reduce_sum(tf.square(B), axis=-1, keepdims=True) # (N, 1)
6 AB = tf.matmul(A, B, transpose_b=True) # (M, N)
7
8 return A_sq - 2 * AB + tf.transpose(B_sq) # (M, N)
9
10dist_sq = pairwise_squared_distance(A, B)
11print(dist_sq) # Same as dist**2
This avoids expand_dims and uses matrix multiplication, which is faster on GPUs.
Pairwise Dot Product
1def pairwise_dot_product(A, B):
2 """Compute dot product between every pair of rows.
3
4 Returns: Tensor of shape (M, N)
5 """
6 return tf.matmul(A, B, transpose_b=True) # (M, D) @ (D, N) = (M, N)
7
8A = tf.constant([[1.0, 0.0], [0.0, 1.0]])
9B = tf.constant([[1.0, 1.0], [2.0, 0.0]])
10
11dots = pairwise_dot_product(A, B)
12print(dots)
13# [[1.0, 2.0],
14# [1.0, 0.0]]
Pairwise Cosine Similarity
1def pairwise_cosine_similarity(A, B):
2 """Compute cosine similarity between all pairs of rows."""
3 A_norm = tf.nn.l2_normalize(A, axis=-1)
4 B_norm = tf.nn.l2_normalize(B, axis=-1)
5 return tf.matmul(A_norm, B_norm, transpose_b=True) # (M, N)
6
7similarity = pairwise_cosine_similarity(A, B)
8print(similarity)
Applying a Custom Function to All Pairs
For arbitrary pairwise operations beyond simple arithmetic:
1def pairwise_apply(A, B, func):
2 """Apply func(a, b) to every pair of rows from A and B.
3
4 Args:
5 A: (M, D)
6 B: (N, D)
7 func: Function that takes two tensors of shape (D,) and returns a scalar
8 """
9 A_expanded = tf.expand_dims(A, axis=1) # (M, 1, D)
10 B_expanded = tf.expand_dims(B, axis=0) # (1, N, D)
11
12 # Tile to create explicit pairs
13 M = tf.shape(A)[0]
14 N = tf.shape(B)[0]
15
16 A_tiled = tf.tile(A_expanded, [1, N, 1]) # (M, N, D)
17 B_tiled = tf.tile(B_expanded, [M, 1, 1]) # (M, N, D)
18
19 return func(A_tiled, B_tiled)
20
21# Example: element-wise maximum, then sum
22result = pairwise_apply(A, B, lambda a, b: tf.reduce_sum(tf.maximum(a, b), axis=-1))
23print(result.shape) # (M, N)
Batch Support
For batched inputs with shape (batch, M, D) and (batch, N, D):
1def batched_pairwise_distance(A, B):
2 """Pairwise distance with batch dimension.
3
4 Args:
5 A: (batch, M, D)
6 B: (batch, N, D)
7
8 Returns: (batch, M, N)
9 """
10 A_expanded = tf.expand_dims(A, axis=2) # (batch, M, 1, D)
11 B_expanded = tf.expand_dims(B, axis=1) # (batch, 1, N, D)
12 diff = A_expanded - B_expanded # (batch, M, N, D)
13 return tf.sqrt(tf.reduce_sum(tf.square(diff), axis=-1))
14
15# Or use matmul for batched dot products
16def batched_pairwise_dot(A, B):
17 return tf.matmul(A, B, transpose_b=True) # (batch, M, N)
Common Pitfalls
Using Python loops instead of broadcasting: Iterating over rows with a for loop is orders of magnitude slower than broadcasting. Always use tf.expand_dims and element-wise operations or tf.matmul for pairwise computations.
Running out of memory with large tensors: Pairwise operations on tensors of shape (M, D) and (N, D) produce (M, N, D) intermediate tensors. For M=N=10000 and D=512, this is 200GB. Use chunked computation or the squared-distance formula with tf.matmul to avoid materializing the full intermediate.
Forgetting transpose_b=True in tf.matmul: tf.matmul(A, B) expects shapes (M, K) and (K, N). To compute A @ B^T for pairwise dot products, pass transpose_b=True instead of manually transposing.
Using tf.sqrt on zero values: tf.sqrt(0.0) has an undefined gradient, causing NaN in backpropagation. Add a small epsilon: tf.sqrt(squared_dist + 1e-8), or use squared distance when the gradient is needed.
Expanding along the wrong axis: tf.expand_dims(A, axis=0) adds a dimension at the front (batch-like), while axis=1 adds it between rows and features. Expanding on the wrong axis produces incorrect broadcasting shapes.
Summary
Use tf.expand_dims to add dimensions for broadcasting: (M, 1, D) and (1, N, D) broadcast to (M, N, D)
For dot products and cosine similarity, use tf.matmul(A, B, transpose_b=True) for efficiency
For squared Euclidean distance, use the expanded formula ||a||^2 - 2*a·b + ||b||^2 to avoid large intermediates
Add 1e-8 before tf.sqrt to prevent NaN gradients at zero distances
Process large pairwise operations in chunks to avoid out-of-memory errors
