テンソルはその内容に基づいてどのように分けるのですか？

Question

私はこのようになりますテンソルあります

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

をそして、私は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. 各要素の第2の指標を合わせるための条件文を書く。.. に私はの行列を返す整合素子をフィルタするためにtf.gatherとtf.not_equal（または同等の機能を）しようとするだろう同じ（0,1）値の入力としての形状と
2. が最初のインデックスによって要素を平均入力（tf.reduce_meanが十分でなければならない）

をフィルタリングするための要素によって乗算