How to update tensors matching dimensionwise vectors

user_6422069 0
Viewed 34

Let there be two 2D tensors, A (m × c) and B (n × c). Each row vector which belongs to B also belongs to A i.e. r. Additionally, row vectors in A are not unique i.e. A may have duplicate rows. However, row vectors in B are unique.

There another pair of tensors P (m × f) and Q (n × f). I am trying to do the following

for i in range(B.shape[0]):
    rv = B[i, :]
    fv = Q[i, :]
    # P[<row indexes of A matching rv>, :] = fv
  1. How to do this correctly?
  2. Is it possible to get rid of the for loop?
torch pytorch numpy python
1 Answers

You can use the following mask:

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
user_90676150
Related