Skip to content

Instantly share code, notes, and snippets.

@mbsariyildiz
Last active March 8, 2024 20:44
Show Gist options
  • Save mbsariyildiz/34cdc26afb630e8cae079048eef91865 to your computer and use it in GitHub Desktop.
Save mbsariyildiz/34cdc26afb630e8cae079048eef91865 to your computer and use it in GitHub Desktop.
Pairwise Euclidean distance computation of elements in 2 tensors, in TensorFlow.
def pairwise_dist (A, B):
"""
Computes pairwise distances between each elements of A and each elements of B.
Args:
A, [m,d] matrix
B, [n,d] matrix
Returns:
D, [m,n] matrix of pairwise distances
"""
with tf.variable_scope('pairwise_dist'):
# squared norms of each row in A and B
na = tf.reduce_sum(tf.square(A), 1)
nb = tf.reduce_sum(tf.square(B), 1)
# na as a row and nb as a co"lumn vectors
na = tf.reshape(na, [-1, 1])
nb = tf.reshape(nb, [1, -1])
# return pairwise euclidead difference matrix
D = tf.sqrt(tf.maximum(na - 2*tf.matmul(A, B, False, True) + nb, 0.0))
return D
@balcilar
Copy link

D=tf.reduce_sum((tf.expand_dims(A, 1)-tf.expand_dims(B, 0))**2,2)
makes more sense

@JayKumarr
Copy link

tf.reduce_sum((tf.expand_dims(A, 1)-tf.expand_dims(B, 0))**2,2)

Tested! Thanks a lot :) Bless you!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment