python - Tensorflow 中的笛卡尔积

Question

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

Answer 1

我在这里假设 a 和 b 都是一维张量。

要获得两者的笛卡尔积，我会结合使用 tf.expand_dims 和 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 和 b 的元素的每个组合都在最后一个维度中表示.