如何更新匹配维度向量的张量

Question

假设有两个二维张量，A (m × c) 和 B (n × c)。属于 B 的每个行向量也属于 A 即。此外，A 中的行向量不是唯一的，即 A 可能有重复的行。但是，B 中的行向量是唯一的。

还有另一对张量 P (m × f) 和 Q (n × f)。我正在尝试执行以下操作

for i in range(B.shape[0]):
    rv = B[i, :]
    fv = Q[i, :]
    # P[<row indexes of A matching rv>, :] = fv

1. 如何正确执行此操作？
2. 是否可以摆脱 for 循环？

Answer 1

您可以使用以下掩码：

for i in range(B.shape[0]):
    rv = B[i]
    fv = Q[i]
    mask = torch.where((A == rv).all(dim=1))[0]
    P[mask] = fv