Filling tensor with zeros after certain index

Dee: <p>Given a 3d tenzor, say: <code>batch x sentence length x embedding dim</code></p>
<pre><code class="lang-auto">a = torch.rand((10, 1000, 96)) 
</code></pre>
<p>and an array(or tensor) of actual lengths for each  sentence</p>
<pre><code class="lang-auto">lengths =  torch .randint(1000,(10,))

outputs tensor([ 370., 502., 652., 859., 545., 964., 566., 576.,1000., 803.])
</code></pre>
<p>How to fill tensor ‘a’ with zeros after certain index along dimension 1 (sentence length) according to tensor ‘lengths’  ?<br>
I want smth like that :</p>
<pre><code class="lang-auto">a[ : , lengths : , : ]  = 0
</code></pre>
<p>One way of doing it (slow if batch size is big enough):</p>
<pre><code class="lang-auto">for i, length in enumerate( lengths ):
    a[ i , length  : , : ]  = 0</code></pre>

zahra: <p>Hi,</p>
<p>Excuse me, if you have a tensor like this:</p>
<pre><code class="lang-auto">tensor([[[ 1.1937, -0.7235, -0.1802, -0.5610],
         [-0.7524,  1.3047, -1.5577,  1.8352],
         [ 1.1573, -1.8952,  0.4175, -0.2085]],

        [[ 0.4069, -0.1069, -0.3838, -0.2991],
         [-0.5824, -0.4965, -0.1542,  1.1482],
         [ 0.5182,  0.5445,  1.1730, -0.2523]]])
</code></pre>
<p>and lengths = torch.tensor([1, 2])<br>
do you want to have sth like this?</p>
<pre><code class="lang-auto">tensor([[[ 1.1937, -0.7235, -0.1802, -0.5610],
         [ 0.0000,  0.0000,  0.0000,  0.0000],
         [ 1.1573, -1.8952,  0.4175, -0.2085]],

        [[ 0.4069, -0.1069, -0.3838, -0.2991],
         [-0.5824, -0.4965, -0.1542,  1.1482],
         [ 0.0000,  0.0000,  0.0000,  0.0000]]])
</code></pre>
<p>or sth like this:</p>
<pre><code class="lang-auto">tensor([[[ 1.4429, -0.8365, -0.2565,  0.2319],
         [ 0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000]],

        [[ 1.7247,  0.2831,  0.0687,  0.6557],
         [ 0.5110,  1.7607,  0.8107,  0.3624],
         [ 0.0000,  0.0000,  0.0000,  0.0000]]])
</code></pre>
<p>Thanks</p>

Dee: <p>I want second option :</p>
<pre><code class="lang-auto">tensor([[[ 1.4429, -0.8365, -0.2565,  0.2319],
         [ 0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000]],

        [[ 1.7247,  0.2831,  0.0687,  0.6557],
         [ 0.5110,  1.7607,  0.8107,  0.3624],
         [ 0.0000,  0.0000,  0.0000,  0.0000]]])
</code></pre>

LeviViana: <p>This works:</p>
<pre><code class="lang-python">&gt;&gt;&gt; import torch
&gt;&gt;&gt; a = torch.rand(2, 3, 4)
&gt;&gt;&gt; print(a)
tensor([[[0.5066, 0.5184, 0.1193, 0.6062],
         [0.4995, 0.1689, 0.6175, 0.7917],
         [0.7996, 0.0225, 0.1145, 0.4249]],

        [[0.0975, 0.2995, 0.5857, 0.0806],
         [0.4922, 0.4778, 0.9133, 0.1418],
         [0.6594, 0.4907, 0.3268, 0.1211]]])
&gt;&gt;&gt; l = torch.randint(3, (2, ))
&gt;&gt;&gt; print(l)
tensor([1, 2])
&gt;&gt;&gt; l += torch.arange(2) * 3
&gt;&gt;&gt; b = torch.arange(6).view(2, 3)
&gt;&gt;&gt; b = b &lt; l.view(2, 1)
&gt;&gt;&gt; a * b.view(2, 3, 1).float()
tensor([[[0.5066, 0.5184, 0.1193, 0.6062],
         [0.0000, 0.0000, 0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000, 0.0000]],

        [[0.0975, 0.2995, 0.5857, 0.0806],
         [0.4922, 0.4778, 0.9133, 0.1418],
         [0.0000, 0.0000, 0.0000, 0.0000]]])
&gt;&gt;&gt; 
</code></pre>
<p>I hope it helps !</p>

zahra: <aside class="quote no-group" data-username="LeviViana" data-post="4" data-topic="53667">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/leviviana/48/9548_2.png" class="avatar"> LeviViana:</div>
<blockquote>
<p>b = b &lt; l.view(2, 1)</p>
</blockquote>
</aside>
<p>What does this command do?</p>

LeviViana: <p>It creates a binary tensor whose values are calculated with respect to the comparison of <code>l</code> and <code>b</code>. The <code>.view(2, 1)</code> is for broadcasting purposes.</p>
<pre><code class="lang-python">&gt;&gt;&gt; import torch
&gt;&gt;&gt; l = torch.randint(3, (2, ))
&gt;&gt;&gt; l # the original form of the lengths
tensor([2, 1])
&gt;&gt;&gt; l += torch.arange(2) * 3
&gt;&gt;&gt; l # the modified form of the lengths for the comparison to work
tensor([2, 4])
&gt;&gt;&gt; b = torch.arange(6).view(2, 3)
&gt;&gt;&gt; b
tensor([[0, 1, 2],
        [3, 4, 5]])
&gt;&gt;&gt; b &lt; l.view(2, 1)
tensor([[1, 1, 0],
        [1, 0, 0]], dtype=torch.uint8)
</code></pre>
<p>The trick is to use <code>torch.arange</code> and then compare it to the indices you want to create the effect of zeroing starting from a given point. I hope it is clear !</p>

Dee: <p>Thanks <a href="https://discuss.pytorch.org/u/LeviViana">LeviViana</a><br>
Nice solution<br>
I also found a bit different approach here:<br>
<aside class="onebox stackexchange">
  <header class="source">
      <a href="https://stackoverflow.com/questions/57548180/filling-torch-tensor-with-zeros-after-certain-index/57550753#57550753" target="_blank" rel="nofollow noopener">stackoverflow.com</a>
  </header>
  <article class="onebox-body">
      <a href="https://stackoverflow.com/users/1714410/shai" target="_blank" rel="nofollow noopener">
    <img alt="Shai" src="https://www.gravatar.com/avatar/a254f486ab076ab1af26162a975e90f9?s=128&amp;d=identicon&amp;r=PG" class="thumbnail onebox-avatar" width="60" height="60">
  </a>
<h4>
  <a href="https://stackoverflow.com/questions/57548180/filling-torch-tensor-with-zeros-after-certain-index/57550753#57550753" target="_blank" rel="nofollow noopener">Filling torch tensor with zeros after certain index</a>
</h4>

<div class="tags">
  <strong>python, nlp, pytorch</strong>
</div>

<div class="date">
  
  answered by
  <a href="https://stackoverflow.com/users/1714410/shai" target="_blank" rel="nofollow noopener">
    Shai
  </a>
  on <a href="https://stackoverflow.com/questions/57548180/filling-torch-tensor-with-zeros-after-certain-index/57550753#57550753" target="_blank" rel="nofollow noopener">05:07AM - 19 Aug 19 UTC</a>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>