Generate groupby index for a tensor

Xuansheng_Wu: <p>Hi guys, I am trying to implement a function that generates unique indexes for a tensor.</p>
<p>The rules are simple:</p>
<ol>
<li>assign a new unique value for each position that value equals to 2.</li>
<li>keep the unique value from the previous positions if the current position is following a 2 started sequence.</li>
<li>ignore value=1 positions which has no 2 a head.</li>
</ol>
<p>Here is an example:<br>
Input: [2, 0, 0, 2, 1, 0, 0, 2, 1, 1, 0, 0, 1, 1, 1]<br>
Expect Output: [1, 0, 0, 2, 2, 0, 0, 3, 3, 3, 0, 0, 0, 0, 0]</p>
<p>I hope the whole pipeline can run within the autograd strategy.</p>
<p>I can successfully select and assign unique values for all values=2 right now, but it is hard for me to select values=1 and lead by 2.</p>
<pre><code class="lang-python">&gt;&gt;&gt; a = tc.Tensor([2., 0, 0, 2, 1, 0, 0, 0, 2, 1, 1, 0, 0, 1, 1, 1])
&gt;&gt;&gt; B = tc.where(a == 2.0, tc.cumsum(a, -1).double(), 0.0)
&gt;&gt;&gt; print(B) 
tensor([2., 0., 0., 4., 0., 0., 0., 0., 7., 0., 0., 0., 0., 0., 0., 0.])
</code></pre>

InnovArul: <aside class="quote no-group" data-username="Xuansheng_Wu" data-post="1" data-topic="151742">
<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/xuansheng_wu/48/49636_2.png" class="avatar"> Xuansheng_Wu:</div>
<blockquote>
<p>assign a new unique value for each position that value equals to 2.</p>
</blockquote>
</aside>
<aside class="quote no-group" data-username="Xuansheng_Wu" data-post="1" data-topic="151742">
<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/xuansheng_wu/48/49636_2.png" class="avatar"> Xuansheng_Wu:</div>
<blockquote>
<p>I hope the whole pipeline can run within the autograd strategy.</p>
</blockquote>
</aside>
<p>I am not sure what do you like to achieve.<br>
AFAIK, the moment you hardcode/assign a value within a tensor, there is a break in the computation graph and the gradient for those elements will be <code>0</code>.</p>

Xuansheng_Wu: <aside class="quote no-group quote-modified" data-username="InnovArul" data-post="2" data-topic="151742" 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/innovarul/48/5282_2.png" class="avatar"> InnovArul:</div>
<blockquote>
<aside class="quote no-group" data-username="Xuansheng_Wu" data-post="1" data-topic="151742">
<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/xuansheng_wu/48/49636_2.png" class="avatar"> Xuansheng_Wu:</div>
<blockquote>
<p>assign a new unique value for each position that value equals to 2.</p>
</blockquote>
</aside>
<aside class="quote no-group" data-username="Xuansheng_Wu" data-post="1" data-topic="151742">
<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/xuansheng_wu/48/49636_2.png" class="avatar"> Xuansheng_Wu:</div>
<blockquote>
<p>I hope the whole pipeline can run within the autograd strategy.</p>
</blockquote>
</aside>
<p>I am not sure what do you like to achieve.<br>
AFAIK, the moment you hardcode/assign a value within a tensor, there is a break in the computation graph and the gradient for those elements will be <code>0</code>.</p>
</blockquote>
</aside>
<p>assume that tensor <code>a</code> is an output of a model and it brings gradients.<br>
I am trying to calculate another tensor <code>b</code> based on <code>a</code> that groups the values in a.</p>
<p>Here is the example:<br>
a = [2, 0, 0, 2, 1, 0, 0, 2, 1, 1, 0, 0, 1, 1, 1]<br>
I need: b = [1, 0, 0, 2, 2, 0, 0, 3, 3, 3, 0, 0, 0, 0, 0]</p>
<p>I hope the algorithm can be done with Pytorch only so that <code>b</code> can carry the gradient from <code>a</code>.</p>

Matias_Vasquez: <p>If I understood your problem correctly, then you could do something like this.</p>
<pre><code class="lang-python">def group_by_index(a):
    b = a.unsqueeze(-2)
    c = b.mT @ b
    x = torch.tril(torch.ones_like(c), -1) + torch.triu(c)
    tmp = (x == 0).cumsum(axis=-1)
    x[tmp&gt;0] = 0
    x, _ = torch.max(x, dim=-2)
    x[x&lt;2] = 0
    x[x&gt;1] = 1
    
    vals = (a==2).cumsum(axis=-1)
    return x * vals


a = torch.tensor([2., 0, 0, 2, 1, 0, 0, 2, 1, 1, 0, 0, 1, 1, 1], requires_grad=True)
x = group_by_index(a)
print(x)
</code></pre>
<pre><code class="lang-python"># Output
tensor([[1., 0., 0., 2., 2., 0., 0., 3., 3., 3., 0., 0., 0., 0., 0.]],
       grad_fn=&lt;MulBackward0&gt;)
</code></pre>
<p>This will also work for tensors with more dims.</p>
<pre><code class="lang-python"># Input
a = torch.tensor([[2., 1, 0, 2, 2],
                  [1, 2, 1, 0, 1],
                  [2, 0, 2, 1, 1]], requires_grad=True)

# Output
tensor([[1., 1., 0., 2., 3.],
        [0., 1., 1., 0., 0.],
        [1., 0., 2., 2., 2.]], grad_fn=&lt;MulBackward0&gt;)
</code></pre>
<hr>
<p>There is definitely room for improvement in this code, but it is just a workaround to get what you want and keep the <code>grad</code>.</p>

Xuansheng_Wu: <p>Thank you Matias!<br>
I am very appreciate to your quick response!<br>
I think this is a perfect solution!</p>

Xuansheng_Wu: <p>Hi Matias,</p>
<p>I tried your solution in my program, but I found that this solution still breaks the backward.<br>
I provide the code I used to check gradient as following:</p>
<pre><code class="lang-python">def group_by_index(a):
    b = a.unsqueeze(-2)
    c = b.mT @ b
    x = torch.tril(torch.ones_like(c), -1) + torch.triu(c)
    tmp = (x == 0).cumsum(axis=-1)
    x[tmp&gt;0] = 0
    x, _ = torch.max(x, dim=-2)
    x[x&lt;2] = 0
    x[x&gt;1] = 1
    
    vals = (a==2).cumsum(axis=-1)
    return x * vals

tags = torch.tensor([2., 0, 0, 2, 1, 0, 0, 2, 1, 1, 0, 0, 1, 1, 1], requires_grad=True)
groups = group_by_index(a)
fake_loss = group.sum() # assume this is final loss of my whole model
fake_loss.backward()
assert tags.grad is not None  # tags.grad currently is None
</code></pre>
<p>I think the reason is that <code>x[tmp &gt; 0]</code> is not differentiable because " index operation won’t be differentiable" from <a href="https://discuss.pytorch.org/t/none-gradients-right-after-parameter-initialization/114604/6" class="inline-onebox">None gradients right after parameter initialization - #6 by ptrblck</a> .</p>
<p>Thank you!</p>

Matias_Vasquez: <p>Hi,</p>
<p>I think you forgot to change the variable that you are giving to the function.</p>
<p><code>a</code> is not defined, but I assume that it did not <code>require_grad</code> in the first place.<br>
I tried with your variable <code>tags</code> and it seems to be working.</p>
<aside class="quote no-group quote-modified" data-username="Xuansheng_Wu" data-post="8" data-topic="151742">
<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/xuansheng_wu/48/49636_2.png" class="avatar"> Xuansheng_Wu:</div>
<blockquote>
<pre><code class="lang-auto">tags = torch.tensor([2., 0, 0, 2, 1, 0, 0, 2, 1, 1, 0, 0, 1, 1, 1], requires_grad=True)
groups = group_by_index(a) # ← Here it should be tags
</code></pre>
</blockquote>
</aside>
<p>However, the gradient seemed to be <code>0</code>, which might have to do with what you mentioned.</p>
<p>I changed it a little bit to avoid these operations</p>
<pre data-code-wrap="python"><code class="lang-python">def group_by_index(a):
    b = a.unsqueeze(-2)
    c = b.mT @ b
    x = torch.tril(torch.ones_like(c), -1) + torch.triu(c)
    tmp = torch.nn.ReLU()((torch.nn.ReLU()((x-1)*-1).cumsum(axis=-1)*-1)+1)
    x = x * tmp
    x, _ = torch.max(x, dim=-2)
    x = x - 1
    x = x / (x+1e-10) # Small value added for nmerical stability
    
    vals = torch.nn.ReLU()(a-1).cumsum(axis=-1)
    return x * vals

tags = torch.tensor([2., 0, 0, 2, 1, 0, 0, 2, 1, 1, 0, 0, 1, 1, 1], requires_grad=True)
groups = group_by_index(tags)
print(groups)
fake_loss = groups.sum() # assume this is final loss of my whole model
fake_loss.backward()
print(tags.grad)
assert tags.grad is not None  # tags.grad currently is None

</code></pre>
<p>Hope this works <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=12" title=":smile:" class="emoji" alt=":smile:" loading="lazy" width="20" height="20"></p>

Matias_Vasquez: <p>This might be a better approach.</p>
<pre><code class="lang-python">def group_by_index(a):
    x = a.unsqueeze(-2).mT @ a.unsqueeze(-2)
    x = torch.tril(torch.ones_like(x), -1) + torch.triu(x)
    mask = ((x-1)*-1).clamp(0, 1).cumsum(axis=-1).clamp(max=1)*-1+1
    x, _ = (x*mask).max(dim=-2)
    x = (x-1).clamp(min=0, max=1)
    
    vals = (a-1).clamp(min=0, max=1).cumsum(axis=-1)
    return x * vals
</code></pre>

Xuansheng_Wu: <p>Thanks to your suggestions! I will look at it in tomorrow.</p>