仅在给定最后一个轴的索引的情况下，如何有效地掩盖张量流中的张量？

Question

假设我有一个形状为(batch_size, a, ... , c, d, e)的张量，其中 a, ... ,c,d,e 是定义的整数。例如(batch_size, 500, 3, 2, 2, 69) 或(batch_size, 2, 2)。

我的问题是针对所有张量的，但让我们继续以 tensor1.get_shape() = (?, 500, 3, 2, 2, 69) 为例

鉴于我有 tensor2 和 tensor2.get_shape() = (?, 500, 3, 2, 2, 14) 包含 tensor1 的最后一个轴的索引，我有 2 个问题：

1) 我想从tensor2 为形状为(?, 500, 3, 2, 2, 69) 的tensor1 构造一个掩码。例如，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 中解决问题并不能解决我的问题，但我仍然想知道。

Answer 1

回答 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])

2 可能没有解决办法，我试试pytorch。