Tensorflow: stack all row pairs from a tensor

Question

Given a tensor t=[[1,2], [3,4]], I need to produce ts=[[1,2,1,2], [1,2,3,4], [3,4,1,2], [3,4,3,4]]. That is, I need to stack together all row pairs. Important: the tensor has dimension [None, 2], ie. the first dimension is variable.

I have tried:

• Using a tf.while_loop to generate a list of indices idx=[[0, 0], [0, 1], [1, 0], [1, 1]], then tf.gather(ts, idx). This works but is messy and I don't know what to do about gradients.
• 2 for loops iterating over tf.unstack(t), adding stacked rows to a buffer, then tf.stack(buffer). This does not work if the first dimension is variable.
• To look for inspiration in broadcasting. For instance, given x=t.expand_dims(t, 0), y=t.expand_dims(t, 1), s=tf.reshape(tf.add(x, y), [-1, 2]) s will be [[2, 4], [4, 6], [4, 6], [6, 8]], ie. the sum of every row combination. But how can I do stacking instead of sum? I've been failing for 2 days :)

Answer 1

Solution with tf.meshgrid() and some reshaping:

import tensorflow as tf
import numpy as np

t = tf.placeholder(tf.int32, [None, 2])
num_rows, size_row = tf.shape(t)[0], tf.shape(t)[1] # actual dynamic dimensions

# Getting pair indices using tf.meshgrid:
idx_range = tf.range(num_rows)
pair_indices = tf.stack(tf.meshgrid(*[idx_range, idx_range]))
pair_indices = tf.transpose(pair_indices, perm=[1, 2, 0])

# Finally gathering the rows accordingly:
res = tf.reshape(tf.gather(t, pair_indices), (-1, size_row * 2))

with tf.Session() as sess:
    print(sess.run(res, feed_dict={t: np.array([[1,2], [3,4], [5,6]])}))
    # [[1 2 1 2]
    #  [3 4 1 2]
    #  [5 6 1 2]
    #  [1 2 3 4]
    #  [3 4 3 4]
    #  [5 6 3 4]
    #  [1 2 5 6]
    #  [3 4 5 6]
    #  [5 6 5 6]]

---

Solution using cartesian product:

import tensorflow as tf
import numpy as np

t = tf.placeholder(tf.int32, [None, 2])
num_rows, size_row = tf.shape(t)[0], tf.shape(t)[1] # actual dynamic dimensions

# Getting pair indices by computing the indices cartesian product:
row_idx = tf.range(num_rows)
row_idx_a = tf.expand_dims(tf.tile(tf.expand_dims(row_idx, 1), [1, num_rows]), 2)
row_idx_b = tf.expand_dims(tf.tile(tf.expand_dims(row_idx, 0), [num_rows, 1]), 2)
pair_indices = tf.concat([row_idx_a, row_idx_b], axis=2)

# Finally gathering the rows accordingly:
res = tf.reshape(tf.gather(t, pair_indices), (-1, size_row * 2))

with tf.Session() as sess:
    print(sess.run(res, feed_dict={t: np.array([[1,2], [3,4], [5,6]])}))
    # [[1 2 1 2]
    #  [1 2 3 4]
    #  [1 2 5 6]
    #  [3 4 1 2]
    #  [3 4 3 4]
    #  [3 4 5 6]
    #  [5 6 1 2]
    #  [5 6 3 4]
    #  [5 6 5 6]]

Answer 2

Can be achieved by:

tf.concat([tf.tile(tf.expand_dims(t,1), [1, tf.shape(t)[0], 1]), tf.tile(tf.expand_dims(t,0), [tf.shape(t)[0], 1, 1])], axis=2)

Detailed steps:

t = tf.placeholder(tf.int32, shape=[None, 2])
#repeat each row of t
d = tf.tile(tf.expand_dims(t,1), [1, tf.shape(t)[0], 1])
#Output:
#[[[1 2] [1 2]]
# [[3 4] [3 4]]]

#repeat the entire input t
e = tf.tile(tf.expand_dims(t,0), [tf.shape(t)[0], 1, 1])
#Output:
#[[[1 2] [3 4]]
# [[1 2] [3 4]]]

#concat
f = tf.concat([d, e], axis=2)

with tf.Session() as sess:
   print(sess.run(f, {t:np.asarray([[1,2],[3,4]])}))  
#Output
#[[[1 2 1 2]
#[1 2 3 4]]
#[[3 4 1 2]
#[3 4 3 4]]]