给定n × n矩阵A，其中A的每一行是[n]的置换，例如，

import torch
n = 100
AA = torch.rand(n, n)
A = torch.argsort(AA, dim=1)

同样给定另一个n × n矩阵P，我们想要构造一个3DTensorQs. t。

Q[i, j, k] = P[A[i, j], k]

在pytorch中有什么有效的方法吗？我知道 torch.gather，但似乎很难直接应用在这里。