Indizierung von Teilsensoren in TensorFlow

Question

Ich möchte einen Ausgangstensor erhalten, indem ich index auf Teilsensoren eines Eingangstensors anwende.

Zum Beispiel in NumPy,

import numpy as np 

input = np.random.random((100,5)) # matrix 
index = np.randint(5, size=(100,)) # vector 
output = data[np.arange(index.shape[0]), index] # vector 

gibt mir gewünschte Ausgabe (Ich mag symbolische Version dieses).

Similiarly in Theano,

import theano.tensor as T 
import theano 

input = T.matrix() # symbolic matrix 
index = T.ivector() # symbolic vector 
output = input[T.arange(index.shape[0]), index] # symbolic vector 

gibt mir eine gewünschte output.

Wie kann ich dies tun in TensorFlow?

import tensorflow as tf 
input = tf.placeholder('float32', [None, 5]) 
index = tf.placeholder('int32', [None]) 
output = ??? 

Im Gegensatz zum Beispiel mit NumPy, die Länge von index (= 1. Dimension der input) nicht festgelegt ist.

Answer 1

Sie können das Slicing tun mit tf.gather_nd:

output = tf.gather_nd(input, tf.stack((tf.range(tf.shape(index)[0]), index), -1))