tensorflow只保留张量top k的值不变，其余值变为0

tensorflow中，只保留张量top k的值不变，其余值变为0，并且对剩余值进行归一化操作
比如：[0.2, 0.2, 0.4, 0.13, 0.07] ，获取top 3的值，但是保持原来向量形状，其余值设为0，结果为 [0.2, 0.2, 0.4, 0, 0]，若要用作权重，最好进行一下归一化操作，使得掩码后的向量和为1，[0.25, 0.25, 0.5, 0, 0]，这样便可以过滤掉一些不重要的信息，在计算加权和时仅仅考虑权重较大的信息

import tensorflow as tf
# python3
# 功能：tensorflow只保留张量top k的值不变，其余值变为0，并且对剩余值进行归一化（最后一维向量除以最后一维总和）
k = 3
a = tf.constant([[[5., 1., 6., 2., 3., 4.],
                  [2.,3.,3.,4.,9.,0.]]])

# a_top: top k 值，  a_top_idx：对应值下标
a_top, a_top_idx = tf.nn.top_k(a, 3, sorted=False)

# 获取第k大值
kth = tf.reduce_min(a_top,axis=2,keepdims=True)

# 大于第k大值的为true，小于为false
top2 = tf.greater_equal(a, kth)

# 映射为掩码，大于第k大值的为1，小于为0
mask= tf.cast(top2,dtype=tf.float32)

# 只保留张量top k的值不变，其余值变为0
v = tf.multiply(a,mask)

# 不为零的数进行归一化操作，即每个数除以该向量的和
sum = tf.reciprocal(tf.reduce_sum(v,axis=2))  # 取倒数
v = tf.transpose(v,[0,2,1])
norms = tf.transpose(tf.multiply(v,sum),[0,2,1])

with tf.Session() as sess:
    sess.run(tf.initialize_all_variables())
    print("top     :\n", sess.run(a_top))
    print("idx     :\n", sess.run(a_top_idx))
    print(sess.run(norms))

    """
    输出：
    top     :
     [[[4. 5. 6.]
      [3. 9. 4.]]]
    idx     :
     [[[5 0 2]
      [1 4 3]]]
      
    [[[0.33333334 0.         0.40000004 0.         0.         0.26666668]
      [0.         0.15789473 0.15789473 0.21052632 0.47368422 0.        ]]]
    """