最後の軸のインデックスを与えられたテンソルのみを効率的にマスクする方法はありますか？

Question

私はテンソル形状を持っているとします。(batch_size, a, ... , c, d, e)ここで、a、...、c、d、eは整数です。例えば、(batch_size, 500, 3, 2, 2, 69)または(batch_size, 2, 2)です。

1）I：

私の質問は、すべてのテンソルのためですが、のは、私はtensor1の最後の軸のtensor2tensor2.get_shape() = (?, 500, 3, 2, 2, 14)で含むインデックスを持っていることを考えるとtensor1.get_shape() = (?, 500, 3, 2, 2, 69)

の例に固執しましょう、私は2つの問題を持っています形状(?, 500, 3, 2, 2, 69)のtensor1のマスクをtensor2から構築したいと考えています。たとえば、tensor2の最後の軸に沿った可能な行は[1,8,3,68,2,4,58,19,20,21,26,48,56,11]ですが、tensor2はtensor1から構成されているため、これらのインデックスは新しい入力で異なります。これらは、保持されなければならない最後の軸のインデックスであるtensor1です。他のすべては隠されなければならない。

2）tensor1の形のマスク(?, 500, 3, 2, 2, 69)を指定した場合、バッチサイズのディメンションを維持しながら望ましくない値をマスクするにはどうすればよいですか？マスクされたテンソルは、形状が(?, 500, 3, 2, 2, 14)である必要があります。

kerasやnumpyの回答は、numpyでそれを行う方法を知っていても問題は解決しませんが、私はまだ知りたいと思います。 1へ

Answer 1

答え： tf.gather_nd(mask, [tf.range(tf.shape(tensor1)[0])[:,None, None, None, None, None],tf.range(tf.shape(tensor1)[1])[:,None, None, None, None],tf.range(tf.shape(tensor1)[2])[:,None, None, None],tf.range(tf.shape(tensor1)[3])[:,None, None],tf.range(tf.shape(tensor1)[4])[:,None],tensor2])

私はpytorchをしようとします。2.への解決策は、おそらくありません。