How to perform cartesian product with Tensorflow?

Question

I'm trying to cross stack two tensors, for example

[0,1,2],[2,3]->[[0,2],[0,3],[1,2],[1,3],[2,2],[2,3]]

Does anybody know which function can do that?

Answer 1

I think this does what you need:

import tensorflow as tf

a = tf.constant([0, 1, 2])
b = tf.constant([2, 3])
c = tf.stack(tf.meshgrid(a, b, indexing='ij'), axis=-1)
c = tf.reshape(c, (-1, 2))
with tf.Session() as sess:
    print(sess.run(c))

Output:

[[0 2]
 [0 3]
 [1 2]
 [1 3]
 [2 2]
 [2 3]]