tensorflow tensor中查找多个元素的索引

import tensorflow as tf

arr=tf.constant([1, 5, 7, 4, 3, 1, 9])
searched_value=tf.constant([1, 5, 9])

index_arr=tf.where(tf.equal(arr,tf.expand_dims(bbb,axis=-1)))
'''
Out:
array([[0, 0],
       [0, 5],
       [1, 1],
       [2, 6]], dtype=int64)
'''
final_index=tf.slice(index_arr,[0,1],[-1,1])#获取索引为1的列
'''
Out: 
array([[0],
       [5],
       [1],
       [6]], dtype=int64)
'''

final_index=tf.reshape(final_index,[-1])#展平tensor,[0,5,1,6]
final_value=tf.gather(arr,final_index)#
'''
Out:
[1,1,5,9]
'''

其中，采用了Broadcasting机制，详情查看https://docs.scipy.org/doc/numpy-1.13.0/user/basics.broadcasting.html

在多维数组中查找多个元素时，如上例子：

arr维度7，searched_value维度3，通过将searched_value在最后插入1维，使searched_value维度变为(3,1)

通过broadcasting机制，实行维度对齐：

arr:                7            1×7            3×7

                      ->补充维度       ->复制

searched_value:   3×1            3×1            3×7


最终：
arr=
[[1, 5, 7, 4, 3, 1, 9],
 [1, 5, 7, 4, 3, 1, 9],
 [1, 5, 7, 4, 3, 1, 9]]
searched_value=
[[1,1,1,1,1,1,1],
 [5,5,5,5,5,5,5],
 [9,9,9,9,9,9,9]]

接着tf.equal(arr,searched_value)->[[True,False,False,False,False,True,False],
                                   [False,True,False,False,False,False,False],
                                    [False,False,False,False,False,False,True]]
tf.where将True对应的索引输出，行索引[0,0,1,2],列索引[0,5,1,6]