Evaluate all pair combinations of rows of two tensors in tensorflow

Question

I am trying to define a custom op in tensorflow, in which at one point I need to construct a matrix (z) that would contain sums of all combinations of pairs of rows of two matrices (x and y). In general, the numbers of rows of x and y are dynamical.

In numpy it's fairly simple:

import numpy as np
from itertools import product

rows_x = 4
rows_y = 2
dim = 2

x = np.arange(dim*rows_x).reshape(rows_x, dim)
y = np.arange(dim*rows_y).reshape(rows_y, dim)

print('x:\n{},\ny:\n{}\n'.format(x, y))

z = np.zeros((rows_x*rows_y, dim))
print('for loop:')
for i, (x_id, y_id) in enumerate(product(range(rows_x), range(rows_y))):
    print('row {}: {} + {}'.format(i, x[x_id, ], y[y_id, ]))
    z[i, ] = x[x_id, ] + y[y_id, ]

print('\nz:\n{}'.format(z))

returns:

x:
[[0 1]
 [2 3]
 [4 5]
 [6 7]],
y:
[[0 1]
 [2 3]]

for loop:
row 0: [0 1] + [0 1]
row 1: [0 1] + [2 3]
row 2: [2 3] + [0 1]
row 3: [2 3] + [2 3]
row 4: [4 5] + [0 1]
row 5: [4 5] + [2 3]
row 6: [6 7] + [0 1]
row 7: [6 7] + [2 3]

z:
[[  0.   2.]
 [  2.   4.]
 [  2.   4.]
 [  4.   6.]
 [  4.   6.]
 [  6.   8.]
 [  6.   8.]
 [  8.  10.]]

However, I haven't got a clue how to implement anything similar in tensorflow.

I was mainly going through SO and the tensorflow API in hopes of finding a function that would yield combinations of elements of two tensors, or a function that would give permutations of elements of a tensor, but to no avail.

Any suggestions are most welcome.

Answer 1

You could simply use the broadcasting ability of tensorflow.

import tensorflow as tf

x = tf.constant([[0, 1],[2, 3],[4, 5],[6, 7]], dtype=tf.float32)
y = tf.constant([[0, 1],[2, 3]], dtype=tf.float32)

x_ = tf.expand_dims(x, 0)
y_ = tf.expand_dims(y, 1)
z = tf.reshape(tf.add(x_, y_), [-1, 2])
# or more succinctly 
z = tf.reshape(x[None] + y[:, None], [-1, 2])

sess = tf.Session()
sess.run(z)