Tensorflow holt Elemente von einem Platzhalter abhängig von einem anderen Platzhalter

Question

Ich habe einige Mühe, diese Operation in tf zu kodieren. Hier ist ein Beispiel, angenommen wir haben einen [n, 2] Platzhalter x und einen [n, 1] Platzhalter y. x = [[1,2], [3,4], [5,6]] y = [1,0,1] für jedes Element i von y möchte ich das entsprechende Element vom ith 2d nehmen Tensor. Im Beispiel sollte die Ausgabe [2,3,6] sein. Ich habe mehrere Techniken ausprobiert, aber ohne Erfolg. Gibt es einen einfachen Weg, das mit Tensorflow zu tun?

Danke

Answer 1

Verwenden Sie entweder tf.gather_nd oder hacken manuell durch tf.stack und tf.where:

import tensorflow as tf 

x = tf.convert_to_tensor([[1, 2], [3, 4], [5, 6]]) 
y = tf.convert_to_tensor([1, 0, 1]) 

with tf.Session() as sess: 
    xx = tf.unstack(x, axis=1) 
    ans = tf.where(tf.equal(y, tf.zeros_like(y)), xx[0], xx[1]) 
    print sess.run(ans) 


with tf.Session() as sess: 
    idx = tf.range(0, limit=3, delta=1, name='arange') 
    idx = tf.stack([idx, y], axis=-1) 
    ans = tf.gather_nd(x, idx) 
    print sess.run(ans)