1

Given tensors x with shape [a, n] and y with shape [b,n] and function f(p,q) operating on tensors of shape [1,n] (or [n]) and returning a scalar value, how can I compute f pair-wise over the batch dimensions of x and y so that my resulting tensor is [a,b,1] (or [a,b])?

I know that this works for operations like multiplication and addition, as described here:Evaluate all pair combinations of rows of two tensors in tensorflow via implicit broadcasting.

How can this be extended to arbitrary functions?

The application is I would like to compute the pair-wise KL divergence for two tensors to match them, so basically a brute force NN computation.

hechth
  • 95
  • 11

1 Answers1

1

One solution (best I could come up using TF) would be to use two map_fn as follows. This will be quite slow as there's no inherent parallelization for map_fn. But you should be able to get some speed-up by tinkering with the parallel_iteration argument.

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))
thushv89
  • 10,865
  • 1
  • 26
  • 39
  • With the function I want to evaluate, this immediately crashes the CUDA solver for some reason, but it's the naive solution. – hechth Dec 13 '19 at 08:06