How to generate variable-length mask?

1

Supposing a is a=tensor([2, 1, 2]), I want to generate a mask based on tensor a .
The mask is like this:

mask = tensor([[ 0,  0,  0,  1,  1],
               [ 0,  0,  0,  0,  1],
               [ 0,  0,  0,  1,  1]])

Is there any simple and convenient method ?

2

The question have solved.
Interesting !

a = torch.tensor([2, 3, 1])
b = torch.arange(10, 0, step=-1).long().expand(3, 4)
a = a.unsqueeze(1).expand_as(c)
mask = b<=a
3

I found this seems to work fine. https://stackoverflow.com/questions/53403306/how-to-batch-convert-sentence-lengths-to-masks-in-pytorch

def length_to_mask(length, max_len=None, dtype=None):
    """length: B.
    return B x max_len.
    If max_len is None, then max of length will be used.
    """
    assert len(length.shape) == 1, 'Length shape should be 1 dimensional.'
    max_len = max_len or length.max().item()
    mask = torch.arange(max_len, device=length.device,
                        dtype=length.dtype).expand(len(length), max_len) < length.unsqueeze(1)
    if dtype is not None:
        mask = torch.as_tensor(mask, dtype=dtype, device=length.device)
    return mask
3 Likes
4

Great link! Prefer the other answer though, personally.

torch.arange(max_len)[None, :] < lens[:, None]

3 Likes