如何从一个内容张量和另一个索引构建张量？

Question

我正在尝试根据其他两个张量的内容组装一个张量，如下所示：

我有一个名为 A 的二维张量，形状为 I * J，还有一个名为 B 的二维张量，形状为 M * N，其元素是 A 的第一维的索引。 我想获得一个形状为 M * N * J 的 3D 张量 C，使得 [0, M) 中的所有 m 和 [0, N) 中的所有 m 都为 C[m,n,:] == A[B[m,n],:]。

我可以使用嵌套的 for 循环来迭代 M 和 N 中的所有索引，在每个索引处为 C 分配正确的值，但是 M 和 N 很大，所以这很慢.我怀疑使用聪明的切片或内置的 pytorch 函数有一些更好、更快的方法，但我不知道它会是什么。它看起来有点像在某个地方使用torch.gather()，但这要求所有张量具有相同的维数。有谁知道这应该怎么做？

编辑：torch.index_select(input, dim, index) 几乎是我想要的，但它在这里不起作用，因为它要求 index 是一维张量，而我的索引张量是二维的。

Answer 1

您可以通过展平第一个维度来实现此目的，让您索引A。需要广播才能恢复最终形状

>>> A[B.flatten(),:].reshape(*B.shape, A.size(-1))

使用A[B.flatten(),:] 进行索引等同于torch.index_select(A, 0, B.flatten())。