Create a mask tensor using index

username: sidwa0
cooked: <p>I’m trying to create *_key_padding_mask for torch.nn.Transformer. So I want to create a mask using the length of sentence data in the batch.<br>
Assuming number of time steps  
<code>t=7</code>
 <code>t=8</code></p>
<p>Edit: changed value of t</p>
<pre><code class="lang-auto">batch_size = 5
idx = torch.tensor([3,4,2,1,6])


mask = torch.tensor([[0,0,0,1,1,1,1,1],
                     [0,0,0,0,1,1,1,1],
                     [0,0,1,1,1,1,1,1],
                     [0,1,1,1,1,1,1,1],
                     [0,0,0,0,0,0,1,1]])
</code></pre>
<p>How do I get <code>mask</code> using <code>idx</code> with vectorized code?</p>

username: Naruto-Sasuke
cooked: <pre><code class="lang-auto">import torch

batch_size = 5

idx = torch.tensor([3,4,1,1,6])

# I changed the mask a little

mask = torch.tensor([[0,0,0,1,1,1,1,1],

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

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

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

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

print(torch.gather(mask, dim=1, index=idx.unsqueeze(-1).expand(mask.size(0), -1)))
</code></pre>

username: sidwa0
cooked: <p>Not getting the expected output.<br>
Mask’s value in my question is the expected output. You already seem to have initialized the mask with the expected output.</p>

username: Naruto-Sasuke
cooked: <p>It seems not easy for current pytorch. Because current operations only accept scalars as the parameters instead of <code>tensor</code>. Something like <code>torch.range(tensor start, tensor end)</code> sovles your problem. You can also refer to <a href="https://github.com/pytorch/nestedtensor" rel="nofollow noopener">https://github.com/pytorch/nestedtensor</a> (which is still experimental)</p>

username: pfloat
cooked: <p>This piece of code may be a good start:</p>
<pre><code class="lang-auto">torch.repeat_interleave(torch.tensor([0, 1, 0, 1]), torch.tensor([3, 8-3 4, 8-4])).view(2, 8)
# gives 
# tensor([[0, 0, 0, 1, 1, 1, 1, 1],
#         [0, 0, 0, 0, 1, 1, 1, 1]])
</code></pre>
<p>We can obtain the one-line code:</p>
<pre><code class="lang-auto">torch.repeat_interleave(torch.tensor([0, 1]*batch_size), torch.stack([idx, 8-idx], dim=1).view(-1)).view(batch_size, 8)
</code></pre>

username: sidwa0
cooked: <p><a class="mention" href="/u/pfloat">@pfloat</a>’s solution with <code>torch.repeat_interleave</code> is quite interesting. However, I found something simple.</p>
<p>Upper triangular matrix</p>
<pre><code class="lang-auto">t = 8
tri = torch.triu( torch.ones(8,8), diagonal=1 )
tri
tensor([[0., 1., 1., 1., 1., 1., 1., 1.],
        [0., 0., 1., 1., 1., 1., 1., 1.],
        [0., 0., 0., 1., 1., 1., 1., 1.],
        [0., 0., 0., 0., 1., 1., 1., 1.],
        [0., 0., 0., 0., 0., 1., 1., 1.],
        [0., 0., 0., 0., 0., 0., 1., 1.],
        [0., 0., 0., 0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0., 0., 0., 0.]])
</code></pre>
<p>Close but if the idx tensor says 3 then index 3 must also be masked.</p>
<pre><code class="lang-auto">tri.fill_diagonal_(1)
tri
tensor([[1., 1., 1., 1., 1., 1., 1., 1.],
        [0., 1., 1., 1., 1., 1., 1., 1.],
        [0., 0., 1., 1., 1., 1., 1., 1.],
        [0., 0., 0., 1., 1., 1., 1., 1.],
        [0., 0., 0., 0., 1., 1., 1., 1.],
        [0., 0., 0., 0., 0., 1., 1., 1.],
        [0., 0., 0., 0., 0., 0., 1., 1.],
        [0., 0., 0., 0., 0., 0., 0., 1.]])
</code></pre>
<p>Now this upper triangular matrix can be referenced to create a mask</p>
<pre><code class="lang-auto">tri[3]
tensor([0., 0., 0., 1., 1., 1., 1., 1.])

idx = torch.tensor([4,1,2,5,3,6])
tri[idx]
tensor([[0., 0., 0., 0., 1., 1., 1., 1.],
        [0., 1., 1., 1., 1., 1., 1., 1.],
        [0., 0., 1., 1., 1., 1., 1., 1.],
        [0., 0., 0., 0., 0., 1., 1., 1.],
        [0., 0., 0., 1., 1., 1., 1., 1.],
        [0., 0., 0., 0., 0., 0., 1., 1.]])
</code></pre>
<p>Since we only need 0 or 1 we can make to dtype bool or byte and this way the method can work for large values of <code>t</code>.</p>

username: pfloat
cooked: <p><a class="mention" href="/u/sidwa0">@sidwa0</a> Nice solution!</p>