After a few hours I finally found how to do this in Tensorflow. My solution works only for Eucledian distance and is pretty verbose. I also do not have a mathematical proof (just a lot of handwaving, which I hope to make more rigorous):
import tensorflow as tf
import numpy as np
from scipy.spatial.distance import cdist
M1, M2, K = 3, 4, 2
# Scipy calculation
a = np.random.rand(M1, K).astype(np.float32)
b = np.random.rand(M2, K).astype(np.float32)
print cdist(a, b, 'euclidean'), '
'
# TF calculation
A = tf.Variable(a)
B = tf.Variable(b)
p1 = tf.matmul(
tf.expand_dims(tf.reduce_sum(tf.square(A), 1), 1),
tf.ones(shape=(1, M2))
)
p2 = tf.transpose(tf.matmul(
tf.reshape(tf.reduce_sum(tf.square(B), 1), shape=[-1, 1]),
tf.ones(shape=(M1, 1)),
transpose_b=True
))
res = tf.sqrt(tf.add(p1, p2) - 2 * tf.matmul(A, B, transpose_b=True))
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
print sess.run(res)
与恶龙缠斗过久,自身亦成为恶龙;凝视深渊过久,深渊将回以凝视…