如何在tensorflow中创建pair组合

Question

我有多个dim张量，我想根据I - dim创建一对组合张量，然后创建两个张量，例如

a=tf.constant([[[1,1],[2,2]],
[[3,3],[4,4]],
[[5,5],[6,6]],shape=(3,2,2)) ,

我根据0 dim(索引为[0,1,2]，所以配对为(0,1),(0,2),(1,2)，所以新张量b的0 dim来自旧索引[0,0,1]，新张量d的0 dim来自旧索引[1,2,2]，最终结果为:

b=tf.constant([[[1,1],[2,2]],
[[1,1],[2,2]],
[[3,3],[4,4]]],shape=(3,2,2))
c=tf.constant([[[3,3],[4,4]],
[[3,3],[4,4]],
[[5,5],[6,6]]],shape=(3,2,2))

Answer 1

使用tf.gather():

import tensorflow as tf
a = tf.constant([[[1,1],[2,2]],
[[3,3],[4,4]],
[[5,5],[6,6]]])
pair = ((0,1),(0,2),(1,2))
pair = tf.convert_to_tensor(pair)
inds = pair[:,0]
b = tf.gather(a, inds)
inds = pair[:,1:]
c = tf.gather(a, inds)