Tensorflowはテンソル内の値のインデックスを取得します

Question

行列とベクトルを考えると、行列の対応する行の値のインデックスを探したいと思います。

m = tf.constant([[0, 2, 1],[2, 0, 1]]) # matrix 
y = tf.constant([1,2]) # values whose indices should be found 

yの最初の値1がmの最初のベクトルのインデックス2にあるため、理想出力は[2,0]です。 yの2番目の値2は、mの2番目のベクトルのインデックス0にあります。

Answer 1

ただ1つの解決策が見つかりました。しかし、よりよいものがあるかどうかは分かりません。

m = tf.constant([[0, 2, 1],[2, 0, 1]]) # matrix 
y = tf.constant([1,2]) # values whose indices should be found 
y = tf.reshape(y, (y.shape[0], 1)) # [[1], [2]] 
cols = tf.where(tf.equal(m, y))[:,-1] # [2,0] 

init = tf.global_variables_initializer() 
with tf.Session() as sess: 
    init.run() 
    print(sess.run(cols))