Create a 2d tensor with varying lengths of one in each row

1

I want to create a mask tensor given a list of lengths. This would mean that there should be k ones and all other zeros for each row in the tensor.

eg:
input :
[2, 3, 5, 1]
output
[1 1 0 0 0
1 1 1 0 0
1 1 1 1 1
1 0 0 0 0]

Is the most efficient way as below or something similar?

seq_lens = torch.tensor([2,3,5,1])
max_len = torch.max(seq_lens)
mask_tensor = torch.Tensor()
for length in seq_lens:
    new_row = torch.cat((torch.ones(length), torch.zeros(max_len-length))).unsqueeze(0)
    mask_tensor = torch.cat((mask_tensor, new_row), 0)
print(mask_tensor)
2

You could use binary masking to achieve this:

seq_lens = torch.tensor([2,3,5,1]).unsqueeze(-1)
max_len = torch.max(seq_lens)

# create tensor of suitable shape and same number of dimensions
range_tensor = torch.arange(max_len).unsqueeze(0)
range_tensor = range_tensor.expand(seq_lens.size(0), range_tensor.size(1))

# until this step, we only created auxiliary tensors (you may already have from previous steps) 
# the real mask tensor is created with binary masking:
mask_tensor = (range_tensor >= seq_lens)
2 Likes