Использование collect_nd для построения матрицы

Question

Я пытаюсь построить тензор ранга-2 tf.float, x, из тензора ранга-3 tf.float, y и тензора ранга-2 tf.int32, z как:

x[i][j] = y[i,z[i][j],j]

Я знаю, что мне нужно использовать tf.gather_nd как:

x = tf.gather_nd(y,indices)

, где

indices[i][j][:] = [i,z[i][j],j]

Однако у меня возникают проблемы с использованием функций тензорного потока для существенного расширения z до более высокого ранга для построения indices.

Я пытаюсь сохранить эти операции в векторизованной форме.

Разумнее ли просто использовать tf.stack as,

indices = tf.stack([ii,z,jj],axis=-1)

, где

ii[i,:] = i

и

jj[:,j] = j

Answer 1

Я думаю, что это делает то, что вам нужно:

import tensorflow as tf
import numpy as np

# Inputs
y = tf.placeholder(tf.float32, [None, None, None])
z = tf.placeholder(tf.int32, [None, None])
# Make first and last indices
y_shape = tf.shape(y)
ii, jj = tf.meshgrid(tf.range(y_shape[0]), tf.range(y_shape[2]), indexing='ij')
# Make full ND index
idx = tf.stack([ii, z, jj], axis=-1)
# Gather result
x = tf.gather_nd(y, idx)
# Test
with tf.Session() as sess:
    # Numbers from 0 to 11 in a (3, 4) matrix
    a = np.arange(12).reshape((3, 4))
    # Make Y with replicas of the matrix multiplied by 1, 10 and 100
    y_val = np.stack([a, a * 10, a * 100], axis=1).astype(np.float32)
    # Z will be a (3, 4) matrix of values 0, 1, 2, 0, 1, 2, ...
    z_val = (a % 3).astype(np.int32)
    # X  should have numbers from 0 to 11 multiplied by 1, 10, 100, 1, 10, 100, ...
    x_val = sess.run(x,  feed_dict={y: y_val, z: z_val}) #, feed_dict={y: y_val, z: z_val})
    print(x_val)

Выход:

[[   0.   10.  200.    3.]
 [  40.  500.    6.   70.]
 [ 800.    9.  100. 1100.]]