Извлечение конкретных элементов из тензора в tensorflow

Question

Я использую tensorflow на Python У меня есть тензор данных формы [?, 5, 37], а тензор IDX формы [?, 5]

Я бы хотел извлечь элементы данных и получить выходную форму таким образом, что [5?]:

output[i][j] = data[i][j][idx[i, j]] for all i in range(?) and j in range(5) 

это выглядит лок на tf.gather_nd() функция ближе всего к моим потребностям, но я не ' я вижу, как использовать его в моем случае ...

Спасибо!

EDIT: Мне удалось сделать это с помощью gather_nd, как показано ниже, но есть ли лучший вариант? (Это кажется немного тяжеловесный)

nRows = tf.shape(length_label)[0] ==> ? 
    nCols = tf.constant(MAX_LENGTH_INPUT + 1, dtype=tf.int32) ==> 5 
    m1 = tf.reshape(tf.tile(tf.range(nCols), [nRows]), 
              shape=[nRows, nCols]) 
    m2 = tf.transpose(tf.reshape(tf.tile(tf.range(nRows), [nCols]), 
              shape=[nCols, nRows])) 
    indices = tf.pack([m2, m1, idx], axis=-1) 
    # indices should be of shape [?, 5, 3] with indices[i,j]==[i,j,idx[i,j]] 
    output = tf.gather_nd(data, indices=indices) 

Answer 1

мне удалось сделать это с gather_nd, как показано ниже

nRows = tf.shape(length_label)[0] # ==> ? 
nCols = tf.constant(MAX_LENGTH_INPUT + 1, dtype=tf.int32) # ==> 5 
m1 = tf.reshape(tf.tile(tf.range(nCols), [nRows]), 
             shape=[nRows, nCols]) 
m2 = tf.transpose(tf.reshape(tf.tile(tf.range(nRows), [nCols]), 
             shape=[nCols, nRows])) 
indices = tf.pack([m2, m1, idx], axis=-1) 
# indices should be of shape [?, 5, 3] with indices[i,j]==[i,j,idx[i,j]] 
output = tf.gather_nd(data, indices=indices)