如何在张量流中成对地评估批量维度上的函数？

Question

给定形状为 [a, n] 的张量 x 和形状为 [b,n] 的 y 和函数 f(p,q) 对形状为 [1,n]（或 [n]）的张量进行运算并返回标量值，如何我可以在 x 和 y 的批量维度上成对计算 f 以便我得到的张量是 [a,b,1]（或 [a,b]）？

我知道这适用于乘法和加法之类的操作，如下所述：通过隐式广播评估张量流中两个张量行的所有对组合。

这如何扩展到任意函数？

应用程序是我想计算两个张量的成对 KL 散度以匹配它们，所以基本上是蛮力 NN 计算。

Answer 1

一种解决方案（我能想到的最好的方法是使用 TF）是使用两个map_fn ，如下所示。 这会很慢，因为map_fn没有固有的并行map_fn 。 但是你应该能够通过修改parallel_iteration参数来获得一些加速。

import tensorflow as tf

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

def f(x,y):
  return x+y

res = tf.map_fn(lambda x: tf.map_fn(lambda y: f(x,y),b), a)

with tf.Session() as sess:
  print(sess.run(res))