Как манипулировать индексами в тензорном потоке?

Question

Я пытаюсь использовать tf.gather_nd для конвертации

'R = tf.eye(3, batch_shape=[4])'

к:

array([[[1., 0., 0.],
        [0., 0., 1.],
        [0., 1., 0.]],

       [[0., 0., 1.],
        [0., 1., 0.],
        [1., 0., 0.]],

       [[0., 1., 0.],
        [0., 0., 1.],
        [1., 0., 0.]],

       [[1., 0., 0.],
        [0., 0., 1.],
        [0., 1., 0.]]], dtype=float32)'

С индексом:

ind = array([[0, 2, 1],
       [2, 1, 0],
       [1, 2, 0],
       [0, 2, 1]], dtype=int32)

Я узнал, могу ли я преобразовать матрицу индекса во что-то вроде:

ind_c = np.array([[[0, 0], [0, 2], [0, 1]], 
               [[1, 2], [1, 1], [1, 0]], 
               [[2, 1], [2, 2], [2, 0]], 
               [[3, 0], [3, 2], [3, 1]]])

collect_nd сделает работу. поэтому мой вопрос:

1. Есть ли лучший способ, чем преобразование индекса ind в ind_c
2. если это единственный способ, как я могу преобразовать ind в ind_c с tenorflow? (Я сделал это сейчас вручную)

Спасибо

Answer 1

Вы можете попробовать следующее:

ind = tf.constant([[0, 2, 1],[2, 1, 0],[1, 2, 0],[0, 2, 1]], dtype=tf.int32)

# Creates the row indices matrix
row = tf.tile(tf.expand_dims(tf.range(tf.shape(ind)[0]), 1), [1, tf.shape(ind)[1]])

# Concat to the ind to form the index matrix
ind_c = tf.concat([tf.expand_dims(row,-1), tf.expand_dims(ind, -1)], axis=2)