如何根据其内容划分张量？

Question

我有一个张量，看起来像这样：

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

我需要计算每个小区的第一元素的平均仅当所述第二元件是等于1。因此，在这个例子中我将选择细胞[3, 1], [1, 1], [5, 1]并计算average(3, 1, 5)=3。 然后，我需要重复这个操作，以0和2结尾的单元格，并总结所有的平均值。

我写这个作为神经网络中损失函数的一个组件，所以我只需要使用TensorFlow函数来完成这个操作。我已经解决了在numpy的这个问题，它看起来像这样：

means_sum = sum([np.mean(np.extract(arr.transpose()[1] == rank, arr.transpose()[0])) for rank in [0, 1, 2]]) 

但我不能找到一个方法来做到这一点在TensorFlow。有任何想法吗？

Answer 1

所以基本上你想

1. 写一个条件语句相匹配的元素.. 我会尝试tf.gather和tf.not_equal（或等同的功能）的第二个指标来筛选返回的矩阵的匹配元素相同的形状在（0,1）值的输入，然后乘以元素来过滤输入
2. 平均由第一索引的元素（tf.reduce_mean应该足够了）