如何根据某些条件为张量中的某些行赋值？

Question

给定张量 left[N,2], right[N,2], mask[1,N]，我想要一种快速方法，根据 right 值的某些条件将行值分配给 left，应用 mask两者，例如：

左=

[[0.9, 0.8],
 [0.3, 0.0],
 [0.6, 0.9],
 [0.7, 0.0],
 [0.6, 0.8],
 [0.6, 0.2],
 [0.6, 0.2]]

面膜=

[ True,  True, False,  True, False, False,  True]

对=

[[ 1.,  3.],
 [ 1.,  5.],
 [ 7.,  0.],
 [11., 13.],
 [17., 19.],
 [21., 1. ],
 [ 1., 13.]]

和，

c1 = [1,3,5]

和

c2 = [11, 13]

。因此，通过

mask

过滤左侧和右侧，如果

right

中的一行同时具有

c1

中的条目或同时存在于

c2

中的条目，则左侧对应的行应为 [0, 1] 如果则不是 [1, 0]。它应用于前向方法，所以我想要一种快速有效的方法来完成它，而不是 for 循环。

所以它应该是这样的，除了这不起作用：

left[mask] = torch.tensor([0, 1]) if (right[mask][0] in c1 and right[mask][1] in c1) or ( right[mask][0] in c2 and right[mask][1] in c2) else torch.tensor([1,0])

Answer 1

left  = torch.tensor([[0.9, 0.8],[0.3, 0.0],[0.6, 0.9],[0.7, 0.0],[0.6, 0.8],[0.6, 0.2],[0.6, 0.2]])
mask  = torch.tensor([ True,  True, False,  True, False, False,  True])
right = torch.tensor([[ 1.,  3.],[ 1.,  5.],[ 7.,  0.],[11., 13.],[17., 19.],[21., 1. ],[ 1., 13.]])
    
c1 = torch.tensor([1,3,5])
c2 = torch.tensor([11, 13])
    
### Function start here
right_in_c1 = right.view(*right.shape, *([1]*len(c1.shape))) == c1.view(*([1]*len(right.shape)), *c1.shape)
right_in_c1 = right_in_c1.view(right.shape[0], -1, torch.prod(torch.tensor(c1.shape))).any(dim= -1).all(dim= 1)
    
right_in_c2 = right.view(*right.shape, *([1]*len(c2.shape))) == c2.view(*([1]*len(right.shape)), *c2.shape)
right_in_c2 = right_in_c2.view(right.shape[0], -1, torch.prod(torch.tensor(c2.shape))).any(dim= -1).all(dim= 1)
    
final_mask = torch.logical_or(right_in_c1, right_in_c2)
    
left[torch.logical_and(mask, final_mask)] = torch.tensor([0, 1]).float()
left[torch.logical_and(mask, ~final_mask)] = torch.tensor([1, 0]).float()
### Function end here
    
left

输出应该是：

tensor([[0.0000, 1.0000],
        [0.0000, 1.0000],
        [0.6000, 0.9000],
        [0.0000, 1.0000],
        [0.6000, 0.8000],
        [0.6000, 0.2000],
        [1.0000, 0.0000]])