Set elements of a matrix prior to a specific index within each row

1

Hi

I have an index vector like:
idx_v = [1,2,4,0,3]
A = [1, 1, 1, 1, 1
1, 1, 1, 1, 1
1, 1, 1, 1, 1
1, 1, 1, 1, 1
1, 1, 1, 1, 1]

I want to set all elements of matrix A located before elements of idx_v vector to zero within each row
like:
At = [0, 1, 1, 1, 1
0, 0, 1, 1, 1
0, 0, 0, 0, 1
1, 1, 1, 1, 1
0, 0, 0, 1, 1]

Is there any way to do it without using For loops?

2

This code should work:

idx_v = torch.tensor([1,2,4,0,3])
A = torch.tensor([[1, 1, 1, 1, 1],
                  [1, 1, 1, 1, 1],
                  [1, 1, 1, 1, 1],
                  [1, 1, 1, 1, 1],
                  [1, 1, 1, 1, 1]])

torch.arange(idx_v)
res = torch.zeros_like(A).scatter_(1, idx_v.unsqueeze(1), 1)
res = torch.cumsum(res, 1)
print(res)
> tensor([[0, 1, 1, 1, 1],
          [0, 0, 1, 1, 1],
          [0, 0, 0, 0, 1],
          [1, 1, 1, 1, 1],
          [0, 0, 0, 1, 1]])
1 Like
3

Thanks it is working