#sequence-mask

Tensor Puzzle 14 keeps values[r, c] when c < length[r] and sets it to 0 otherwise.

def arange(i):
    return torch.tensor(range(i))

def where(q, a, b):
    return (q * a) + (~q) * b

values = torch.ones(3, 4, dtype=torch.long)
length = torch.tensor([1, 3, 0])

Which line computes it correctly?