Tensorflow 中的笛卡尔积

有没有像 itertools.product 这样在 Tensorflow 中进行笛卡尔积的简单方法？我想获得两个张量元素的组合（a and b），在Python中可以通过itertools作为list(product(a, b))。我正在寻找 Tensorflow 的替代方案。

我在这里假设两者a and b是一维张量。

为了获得两者的笛卡尔积，我将使用以下组合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) 

你最终得到一个 len(a) * len(b) * 2 张量，其中元素的每个组合a and b表示在最后一个维度中。