Is there any easy way to do cartesian product in Tensorflow like itertools.product? I want to get combination of elements of two tensors (a
and b
), in Python it is possible via itertools as list(product(a, b))
. I am looking for an alternative in Tensorflow.
I'm going to assume here that both a
and b
are 1-D tensors.
To get the cartesian product of the two, I would use a combination of tf.expand_dims
and tf.tile
:
a = tf.constant([1,2,3])
b = tf.constant([4,5,6,7])
tile_a = tf.tile(tf.expand_dims(a, 1), [1, tf.shape(b)[0]])
tile_a = tf.expand_dims(tile_a, 2)
tile_b = tf.tile(tf.expand_dims(b, 0), [tf.shape(a)[0], 1])
tile_b = tf.expand_dims(tile_b, 2)
cartesian_product = tf.concat([tile_a, tile_b], axis=2)
cart = tf.Session().run(cartesian_product)
print(cart.shape)
print(cart)
You end up with a len(a) * len(b) * 2 tensor where each combination of the elements of a
and b
is represented in the last dimension.
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]]
credit to jdehesa: link
A shorter solution to the same, using tf.add()
for broadcasting (tested):
import tensorflow as tf
a = tf.constant([1,2,3])
b = tf.constant([4,5,6,7])
a, b = a[ None, :, None ], b[ :, None, None ]
cartesian_product = tf.concat( [ a + tf.zeros_like( b ),
tf.zeros_like( a ) + b ], axis = 2 )
with tf.Session() as sess:
print( sess.run( cartesian_product ) )
will output:
[[[1 4]
[2 4]
[3 4]][[1 5]
[2 5]
[3 5]][[1 6]
[2 6]
[3 6]][[1 7]
[2 7]
[3 7]]]
A more succinct version of Sunreef's answer uses tf.stack
instead of tf.concat
a = tf.constant([1,2,3])
b = tf.constant([4,5,6,7])
tile_a = tf.tile(tf.expand_dims(a, 1), [1, tf.shape(b)[0]])
tile_b = tf.tile(tf.expand_dims(b, 0), [tf.shape(a)[0], 1])
ans = tf.stack([tile_a, tile_b], -1)
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With