来自one-hot编码的Tensorflow掩码

Question

我有examples = tf.placeholder(tf.int32, [batch_size])形式的标签，其中每个示例都是范围0:ohe_size中的一个int。

我的输出是带有形状[batch_size, ohe_size]的softmax概率分布的形式

我正在尝试创建一个掩码，它将只给出每个示例的概率分布。例如：

probs = [[0.1, 0.6, 0.3]
         [0.2, 0.1, 0.7]
         [0.9, 0.1, 0.0]]
examples = [2, 2, 0]

some_mask_func(probs, example) # <- Need this function    
> [0.3, 0.7, 0.9]