How to index 4D tensors with 2D tensors?

username: Jeonsam
cooked: <p>Hi!<br>
I have a 4D tensor A with shape (B, X, Y, Z) and a tensor B with shape (N, 4), each of whose rows stores the index of one element of tensor A.<br>
For example:</p>
<pre><code class="lang-auto">B = torch.Tensor([[0, 0, 0, 0], [0, 0, 0, 1]])
# I want to access tensor A:
A[0, 0, 0, 0] = 1
A[0, 0, 0, 1] = 1
</code></pre>
<p>How can I access tensor A using tensor B without for loop? I find it difficult to use scatter, index_select, and index_put.</p>
<p>Thank you!</p>

username: carperbr
cooked: <p>If I’m understanding you correctly, you should be able to use a long tensor.</p>
<pre><code class="lang-auto">b = torch.Tensor([[0, 0, 0, 0], [0, 0, 0, 1]]).type(torch.long)
a[b] = 1
</code></pre>

username: Jeonsam
cooked: <p>Thank you for your reply.<br>
I tried your method but get:</p>
<pre><code class="lang-auto">RuntimeError: CUDA out of memory
</code></pre>
<p>I think your solution will select an entire row rather than one single element.</p>
<pre><code class="lang-auto">a = torch.arange(8).view(2, 4)
a
tensor([[0, 1, 2, 3],
        [4, 5, 6, 7]])

b = torch.tensor([[0, 0, 0, 0], [0, 0, 0, 1]]).type(torch.long)
a[b]
tensor([[[0, 1, 2, 3],
         [0, 1, 2, 3],
         [0, 1, 2, 3],
         [0, 1, 2, 3]],

        [[0, 1, 2, 3],
         [0, 1, 2, 3],
         [0, 1, 2, 3],
         [4, 5, 6, 7]]])
</code></pre>

username: carperbr
cooked: <p>Ah yeah whoops, not sure if there is any way of doing that via indexing unfortunately</p>

username: Matias_Vasquez
cooked: <p>You could do something like this.</p>
<pre><code class="lang-python">a = torch.arange(8).view(1, 1, 2, 4)
b = torch.tensor([[0, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 3]])

print(a[b[:, 0], b[:, 1], b[:, 2], b[:, 3]])
</code></pre>
<pre><code class="lang-python"># Output:
tensor([0, 1, 7])
</code></pre>

username: Jeonsam
cooked: <p>Thank you. Your method is simple and cool. <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"> I figure out a method using index_put_:</p>
<pre><code class="lang-auto">a = torch.arange(8).view(1, 1, 2, 4)
b = torch.tensor([[0, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 3]])
inds = torch.chunk(b, 4, dim=1)
a.index_put_(inds, torch.tensor(0))
</code></pre>
<p>Not sure which one runs faster. I test them with pytorch1.9:</p>
<pre><code class="lang-auto">import torch
import time

a = torch.arange(468 * 468, device='cuda:0').view(1, 1, 468, 468)
b = torch.randint(0, 468, (20000, 4), device='cuda:0')
b[:, 0] = b[:, 1] = 0

t1 = time.time()
for i in range(10000):
    a[b[:, 0], b[:, 1], b[:, 2], b[:, 3]] = 1
print(f"Time: {time.time() - t1}") 

t2 = time.time()
for i in range(10000):
    inds = torch.chunk(b, 4, dim=1)
    a.index_put_(inds, torch.tensor(1).to(b.device))
print(f"Time: {time.time() - t2}")
</code></pre>
<pre><code class="lang-auto"># Output:
Time: 0.5770895481109619
Time: 0.7494716644287109
</code></pre>
<p>Using index_put_ is slower.<br>
And index_put_ shows lower speed even when not considering the overhead of torch.chunk. (0.41 vs 0.44)</p>

username: Jeonsam
cooked: <p>I am sorry. Index_put_ seems to be much faster.<br>
Torch.chunk and torch.tensor(1).to(b.device) cause the slowness.</p>
<pre><code class="lang-auto">import torch
import time

a = torch.arange(468 * 468, device='cuda:0').view(1, 1, 468, 468)
b = torch.randint(0, 468, (20000, 4), device='cuda:0')
b[:, 0] = b[:, 1] = 0

t1 = time.time()
for i in range(10000):
    a[b[:, 0], b[:, 1], b[:, 2], b[:, 3]] = 1
print(f"Time: {time.time() - t1}") 

inds = torch.chunk(b, 4, dim=1)
value = torch.tensor(1).to(b.device)
t2 = time.time()
for i in range(10000):
    a.index_put_(inds, value)
print(f"Time: {time.time() - t2}")
</code></pre>
<pre><code class="lang-auto"># Output:
Time: 0.38867831230163574
Time: 0.13382434844970703
</code></pre>