Removing Zeros (between non-zero values) and maintaining the Tensor dimensions

username: Charamba
cooked: <p>Hi Guys,</p>
<p>I would like to remove zero values of a tensor and “join” the non-zero values in each row of a tensor in format [B, C, H, W]. A naive way would to do <code>out_x = x[x!=0]</code>, this approach is bad because would destruct the Tensor dimensions. <strong>For resume, I would like to transform an input tensor like this:</strong></p>
<pre><code class="lang-auto">in_x = torch.tensor([[[[0.,   0.,   2.,  20., 250.,   0.,   0.,   0.,   0.,   0., 250.,
                      20.,   2.,   0.,   0.,   0.],
                      [0.,   0.,   2.,  20.,   0.,   0.,  20., 250.,   0.,   0., 250.,
                      20.,  20.,   2.,   0.,   0.],
                      [0.,   0.,   0.,   0.,   0.,   0.,   2.,  20., 250.,   0., 250.,
                      20.,  20.,   2.,   0.,   0.],
                      [0.,   2.,  20.,   0.,  20., 250.,   0., 250.,  20.,   0.,  20.,
                       2.,   0.,   0.,   0.,   0.]]]])
</code></pre>
<p><strong>In an output tensor like this:</strong></p>
<pre><code class="lang-auto">out_x = torch.tensor([[[[  2.,  20., 250., 250.,  20.,   2.,   0.,   0.,   0.,   0.,   0.,
                          0.,   0.,   0.,   0.,   0.],
                       [  2.,  20.,  20., 250., 250., 20.,  20.,   2.,    0.,   0.,   0.,
                          0.,   0.,   0.,   0.,   0.],
                       [  2.,  20., 250., 250.,  20., 20.,   2.,   0.,   0.,   0.,   0.,
                          0.,   0.,   0.,   0.,   0.],
                       [  2.,  20.,  20., 250., 250.,  20.,  20.,   2., 0.,   0.,   0.,
                          0.,   0.,   0.,   0.,   0.]]]])
</code></pre>
<p><strong>Note that both Tensors have the same shape: (1,1,4,16)</strong></p>
<p>The solution must not contain For Loops or anything else that degrades performance.</p>
<p>Thanks.</p>

username: KFrank
cooked: <p>Hi Luiz!</p>
<aside class="quote no-group quote-modified" data-username="Charamba" data-post="1" data-topic="140062" data-full="true">
<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/charamba/48/45194_2.png" class="avatar"> Charamba:</div>
<blockquote>
<p>Hi Guys,</p>
<p>I would like to remove zero values of a tensor<br>
…</p>
<p><strong>In an output tensor like this:</strong></p>
<pre><code class="lang-auto">out_x = torch.tensor([[[[  2.,  20., 250., 250.,  20.,   2.,   0.,   0.,   0.,   0.,   0.,
                          0.,   0.,   0.,   0.,   0.],
                       [  2.,  20.,  20., 250., 250., 20.,  20.,   2.,    0.,   0.,   0.,
                          0.,   0.,   0.,   0.,   0.],
                       [  2.,  20., 250., 250.,  20., 20.,   2.,   0.,   0.,   0.,   0.,
                          0.,   0.,   0.,   0.,   0.],
                       [  2.,  20.,  20., 250., 250.,  20.,  20.,   2., 0.,   0.,   0.,
                          0.,   0.,   0.,   0.,   0.]]]])
</code></pre>
</blockquote>
</aside>
<p>At the cost of  <code>n log (n)</code> time complexity, you can use <code>argsort()</code>* and<br>
then use <code>gather()</code> to index back into your input tensor:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.9.0'
&gt;&gt;&gt; in_x = torch.tensor([[[[0.,   0.,   2.,  20., 250.,   0.,   0.,   0.,   0.,   0., 250.,
...                       20.,   2.,   0.,   0.,   0.],
...                       [0.,   0.,   2.,  20.,   0.,   0.,  20., 250.,   0.,   0., 250.,
...                       20.,  20.,   2.,   0.,   0.],
...                       [0.,   0.,   0.,   0.,   0.,   0.,   2.,  20., 250.,   0., 250.,
...                       20.,  20.,   2.,   0.,   0.],
...                       [0.,   2.,  20.,   0.,  20., 250.,   0., 250.,  20.,   0.,  20.,
...                        2.,   0.,   0.,   0.,   0.]]]])
&gt;&gt;&gt;
&gt;&gt;&gt; in_x.gather (3, (in_x == 0.0).sort (dim = 3, stable = True)[1])
tensor([[[[  2.,  20., 250., 250.,  20.,   2.,   0.,   0.,   0.,   0.,   0.,
             0.,   0.,   0.,   0.,   0.],
          [  2.,  20.,  20., 250., 250.,  20.,  20.,   2.,   0.,   0.,   0.,
             0.,   0.,   0.,   0.,   0.],
          [  2.,  20., 250., 250.,  20.,  20.,   2.,   0.,   0.,   0.,   0.,
             0.,   0.,   0.,   0.,   0.],
          [  2.,  20.,  20., 250., 250.,  20.,  20.,   2.,   0.,   0.,   0.,
             0.,   0.,   0.,   0.,   0.]]]])
</code></pre>
<p>*)  Except you have to use <code>sort()</code> because <code>argsort()</code> doesn’t support<br>
<code>stable = True</code>).</p>
<p>Best.</p>
<p>K. Frank</p>

username: Charamba
cooked: <aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="140062">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p><code>in_x.gather (3, (in_x == 0.0).sort (dim = 3, stable = True)[1])</code></p>
</blockquote>
</aside>
<p>Thank you very much K. Frank, this works fine. Have a Merry Christmas and a Happy New Year!</p>