Sort index according to label permutation

MehdiZouitine: <p>Hello !</p>
<p>I have a batch of tensor that represent a permutation of label called label_perm.<br>
I have another batch tensor called node_label that give the label of each node.<br>
I want to have a batch tensor of node index called res that follow the order of the label_perm :</p>
<p>for example :</p>
<p>label_perm = [[0,1,2]] <span class="hashtag">#the</span> label 0 is the first, the label 1 come after and finally the label 2<br>
node_label = [[1,2,1,0,1,1,2]] # each node has a label<br>
res = [[3,0,2,4,5,1,6]]<br>
The first label in label_perm is 0 and only the node 3 has label 0. Then the second label is 1 so the node 0,2,4,5 have label 1 and so on …<br>
Do you know how to do this in torch with an arbitrary batch size ?</p>
<p>Thanks you very much ! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

tataganesh: <p>This isn’t a clean solution, but it might at least help you get started.</p>
<p><strong>Defining Tensors</strong></p>
<pre><code class="lang-auto">label_perm = torch.tensor([[0,1,2]])
node_label = torch.tensor([[1,2,1,0,1,1,2]])
</code></pre>
<p><strong>Broadcasting</strong></p>
<p>Assume that the shape of <code>node_label</code> is <code>(1, m)</code> and the shape of <code>label_perm</code> is <code>(1, n)</code>.  We want to make the shape of <code>label_perm</code> to be <code>(n,1)</code> (using <code>reshape</code>) to allow broadcasting. Broadcasting enables us to create a tensor <code>value_present</code> of shape <code>n x m</code>, where <code>result[n,m]</code> is <code>True</code> if <code>label_perm[n] == node_label[m]</code>.</p>
<pre><code class="lang-auto">value_present = (node_label == label_perm.reshape(-1, 1))
"""
value_present =
 tensor([[False, False, False,  True, False, False, False],
        [ True, False,  True, False,  True,  True, False],
        [False,  True, False, False, False, False,  True]])
"""
</code></pre>
<p>Now, get the indices of all <code>True</code> values in <code>value_present</code> and store it in <code>indices</code>. The second column of <code>indices</code> should yield the final tensor.</p>
<pre><code class="lang-auto">indices = value_present.nonzero()
result = indices[:, 1]
"""
indices = 
  tensor([[0, 3],
          [1, 0],
          [1, 2],
          [1, 4],
          [1, 5],
          [2, 1],
          [2, 6]])
result = tensor([3, 0, 2, 4, 5, 1, 6])
"""
</code></pre>
<p>You can modify this to work for arbitrary batch sizes.</p>

MehdiZouitine: <aside class="quote no-group" data-username="MehdiZouitine" data-post="1" data-topic="171501">
<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/mehdizouitine/48/15365_2.png" class="avatar"> MehdiZouitine:</div>
<blockquote>
<p>I have a batch of tensor that represent a permutation of label called label_perm.<br>
I have another batch tensor called node_label that give the label of each node.<br>
I want to have a batch tensor of node index called res that follow the order of the label_perm :</p>
<p>for example :</p>
<p>label_perm = [[0,1,2]] <span class="hashtag-raw">#the</span> label 0 is the first, the label 1 come after and finally the label 2<br>
node_label = [[1,2,1,0,1,1,2]] # each node has a label<br>
res = [[3,0,2,4,5,1,6]]<br>
The first label in label_perm is 0 and only the node 3 has label 0. Then the second label is 1 so the node 0,2,4,5 have label 1 and so on …<br>
Do you know how to do this in torch with an arbitrary batch size ?</p>
</blockquote>
</aside>
<p>Thank you very much for you answer, it really do the job !</p>
<p>I found another solution by myself using one_hot_encoding and bmm <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<pre><code class="lang-auto">label_perm = torch.tensor([[0,1,2],[0,1,2]]) 
node_label = torch.tensor([[1,2,1,0,1,1,2],[1,2,1,1,1,1,2]])
res = torch.tensor([[3,0,2,4,5,1,6],[0,2,3,4,5,1,6]])
a = F.one_hot(label_perm,num_classes=3)
b=F.one_hot(node_label,num_classes=3)
z = torch.bmm(a,b.transpose(-1,-2))
k = torch.argwhere(z&gt;0)
perm = k[:,-1].view(2,-1)

</code></pre>
<p>Here an example <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>