Create masking tensor using lower and upper boundary tensors

1

I want to create a masking tensor of shape [N, K] that contains ones in between different ranges of K for each N and zeros outside. To create this, I have two tensors with shape [N, 1] that contain the lower and upper index for each example that I should give “one”. In other words:

Assuming N=3 and K=5, as well as the lower limit tensor is [[2], [0], [4]] and the upper limit tensor is [[3], [2], [4]] then, the masking tensor should look like this:

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

Is there a way of creating this masking tensor with the provided Pytorch functions?

1 Like
2

Hi,
You can use torch.ge (greater or equal) and torch.le (lower or equal) with :

lower = torch.LongTensor([2, 0, 4]).unsqueeze(dim=-1)
upper = torch.LongTensor([3, 2, 4]).unsqueeze(dim=-1)

idx = torch.arange(5)
mask = idx.ge(lower) & idx.le(upper)

It will give you a tensor with True and False
Use mask.long() if you need ones and zeros

1 Like