Is there any API can work for this?

1

I have a tensor A of shape m*n and a long tensor B of shape m*1, each element in B is in [0, n-1] indicating the last unmasked element in each row of A, and then I can return a mask of m*n for A based on B.

For example, the shape of A is 4*3, and B is [[2], [1], [0], [1]], then I want to return a mask as follows,

[[1, 1, 1],
[1, 1, 0],
[1, 0, 0],
[1, 1, 0]]

Is there any similar API provided by PyTorch?

2

Hi,

I don’t think there is I’m afraid.
The best way I can think of to do this is:

out = torch.zeros_like(A)
n = A.size(1)
out.scatter_(1, n - 1 - B, 1) # Work in reverse order
out = out.cumsum(dim=1)
out.flip(1)