Selecting and concatenating rows

cajoek: <p>Hi! I have two tensors:</p>
<ul>
<li>
<strong>values</strong>, FloatTensor, shape: [32, 4096, 3], value ∈ ℝ</li>
<li>
<strong>indices</strong>, LongTensor, shape: [32, 4096, 4], index ∈ {0…4095}</li>
</ul>
<p>The indices point to dim 1 in values. How can I create a new tensor of shape: [32, 4096, 12] by selecting and concatenating triplets from the values tensor.</p>
<p><strong>Edit:</strong><br>
I came up with this but it doesn’t seem very efficient:</p>
<pre><code class="lang-auto">x =torch.cat(
    [torch.gather(mesh, 1, index=indices[:, :, i].unsqueeze(-1).expand(-1, -1, 3)) for i in range(3)], dim=2
)
</code></pre>

Matias_Vasquez: <p>You can do something like this.</p>
<pre><code class="lang-python">cols= torch.tensor([0, 4, 8, 9, 1, 5, 6, 10, 2, 3, 7, 11])
y = torch.gather(mesh.repeat(1, 1, 4), dim=1, index=indices.repeat(1, 1, 3))[:, :, cols]
</code></pre>
<p>(I do not know why it has this weird format, but in order to get the same result as yours I had to shuffle the columns like that)<br>
I mean, it does make sense that they are grouped into:</p>
<ul>
<li>0, 4, 8</li>
<li>1, 5, 9</li>
<li>2, 6, 10</li>
<li>3, 7, 11</li>
</ul>
<p>But why they are always permuting, I cannot answer that.</p>
<p>If you do not care that the results are permuted, then you can use this ↓. The values are the same, just not in the same column order.</p>
<pre><code class="lang-python">y = torch.gather(mesh.repeat(1, 1, 4), dim=1, index=indices.repeat(1, 1, 3))
</code></pre>
<h2>
<a name="test-1" class="anchor" href="#test-1"></a>Test</h2>
<p>Just to make sure that the result is the same.</p>
<pre><code class="lang-python">mesh = torch.rand(32, 4096, 3)
indices = torch.randint(0, 4095, (32, 4096, 4))

x = torch.cat(
    [torch.gather(mesh, 1, index=indices[:, :, i].unsqueeze(-1).expand(-1, -1, 3)) for i in range(4)], dim=2
)

cols = torch.tensor([0, 4, 8, 9, 1, 5, 6, 10, 2, 3, 7, 11])
y = torch.gather(mesh.repeat(1, 1, 4), dim=1, index=indices.repeat(1, 1, 3))[:, :, cols]

print(x.shape)
print(y.shape)
print(torch.all(x==y))
</code></pre>
<h2>
<a name="result-2" class="anchor" href="#result-2"></a>Result</h2>
<pre><code class="lang-python">torch.Size([32, 4096, 12])
torch.Size([32, 4096, 12])
tensor(True)
</code></pre>

cajoek: <p>Thank you <a class="mention" href="/u/matias_vasquez">@Matias_Vasquez</a>!</p>
<p>Pytorch indexing (and numpy if it is the same) is a complete mystery to me. I also found this solution:</p>
<pre><code class="lang-auto">mesh = torch.rand(32, 4096, 3)
indices = torch.randint(0, 4095, (32, 4096, 4))

x = torch.cat([torch.gather(mesh, 1, index=indices[:, :, i].unsqueeze(-1).expand(-1, -1, 3)) for i in range(4)], dim=2)

cols = torch.tensor([0, 4, 8, 9, 1, 5, 6, 10, 2, 3, 7, 11])
y = torch.gather(mesh.repeat(1, 1, 4), dim=1, index=indices.repeat(1, 1, 3))[:, :, cols]

# Another solution
i = indices.view(indices.size(0), indices.size(1) * indices.size(2))
z = mesh[torch.arange(mesh.size(0)).unsqueeze(-1), i].view(mesh.size(0), mesh.size(1), -1)
##

print(x.shape)
print(y.shape)
print(z.shape)
print(torch.all(x == y))
print(torch.all(x == z))
</code></pre>
<p>Results:</p>
<pre><code class="lang-auto">torch.Size([32, 4096, 12])
torch.Size([32, 4096, 12])
torch.Size([32, 4096, 12])
tensor(True)
tensor(True)
</code></pre>
<p>But I’m not sure which is better, I don’t understand any of them</p>