对 Tensorflow 中两个张量中的每对元素应用函数

Question

我有两个 rank-2 Tensors，在第二个维度上大小相等，但在第一个维度上不相等。例如，形状为[a, n] 的张量A 和形状为[b, n] 的张量B。它们可以看作是两个包含长度为n的向量的数组。

我有一个函数f，它接受两个输入，每个输入一个形状为[n] 的张量，并返回一个标量。我想将此函数应用于A 和B 中的每对向量，结果是形状为[a, b] 的张量C，这样，对于C 中的每个位置(i, j)，C[i, j] = f(A[i], B[j]) .

如果这些只是常规的 Numpy 数组，我可以使用以下代码完成此操作：

# Assume a, b, and n are integers, and A and B are Numpy arrays
C = numpy.zeros((a, b))
for i in range(0, a):
    for j in range(0, b):
        C[i, j] = f(A[i], B[j])
return C

如果这可以通过f 简单地将A 和B 作为输入并返回C 的方式来完成，那将是首选的解决方案，以便一切都作为适当的张量操作发生，所以这一切都可以通过 Tensorflow 正确并行化。只要最终结果相同。

我找到了一个solution to this problem，专门用于f 计算每对向量之间的欧几里得距离。我想将此扩展到其他函数，例如余弦距离或曼哈顿 (L1) 距离。

Answer 1

a = tf.random_normal([10,5])
b = tf.random_normal([20,5])

我将首先像这样重新定位两个数组：

a = a[:,tf.newaxis,:]
b = b[tf.newaxis,:,:]

现在形状是 [a,1,n] 和 [1,b,n]，所以我们可以广播一个减法来计算每对的 delta：

delta = (a-b)

它的形状为 [a,b,n]。

现在欧几里得距离是直截了当的。 （axis=-1 在最后一个轴上求和）：

distance = tf.reduce_sum(delta**2,axis = -1)**0.5

你就完成了：

print(distance)
<tf.Tensor 'pow_3:0' shape=(10, 20) dtype=float32>