Premium problem101. Padding Mask From Lengths

Medium Locked

Given lengths of shape (B,) and an integer max_len, return a (B, max_len) boolean tensor that is True where a position is within that row's length and False where it is padding.

So lengths = [3, 1, 4] with max_len = 5 gives a first row of [True, True, True, False, False].

No Python loops -- compare a row of positions against the lengths and let broadcasting do the work.

Input

lengths = tensor([3, 1, 4])
max_len = 5

Output

tensor([[ True,  True,  True, False, False],
        [ True, False, False, False, False],
        [ True,  True,  True,  True, False]])

Premium problem

This one's part of Premium. Unlock the full PyTorch track plus every other premium problem on the site.

Implement solve(...)