How to vectorize this pytorch code over (at least) the batch dimension?

Natasha_667: <p>Please help me to vectorize this code on the batch dimension at least:</p>
<pre><code class="lang-auto">A = torch.zeros((3, 3, 3), dtype = torch.float)
X = torch.tensor([[0, 1, 2, 0, 1, 0], [1, 0, 0, 2, 1, 1], [0, 0, 2, 2, 1, 1]])
for a, x in zip(A, X):
    for i, j in zip(x, x[1:]):
        a[i, j] = 1
</code></pre>
<p>Thanks! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

albanD: <p>Hi,</p>
<p>When you say the batch dimension, is it the first dimension in your example? Or another dimension to add to your example?<br>
Also the formula that you want is the following?<br>
<code>forall i in A.size(0), forall j in x.size(0) - 1: A[i, x[i, j], x[i, j+1]] = 1</code><br>
Meaning that every pair of consecutive values in x are the the indices where you should put a 1 ?</p>

Natasha_667: <p>Hi! That’s right! The batch dimension is the first.</p>

albanD: <p>Just to be sure before writing code for this (to make sure I don’t write code that does not match what you want), what is the rational for the consecutive numbers in x being used as coordinate? Why each number  is first an index on the 2nd dimension then an index on the 1st dimension?<br>
It would fell more natural to have half being indices in the 1st dimension and the other half being indices in the 2nd dimension. Or at least have them used only once, not twice.</p>

Natasha_667: <p>Consider X[0] : [0, 1, 2, 0, 1, 0]. What I want to build an adjacency matrix such that:</p>
<blockquote>
<p>A[0, 1] = 1<br>
A[1, 2] = 1<br>
A[2, 0] = 1<br>
A[0, 1] = 1<br>
A[1, 0] = 1</p>
</blockquote>
<p>The code I shared works, however it’s very slow since I’m looping over each element in the batch. I’d like to vectorize it as much as possible. Thanks! <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>

albanD: <p>There you go, this will be much faster <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
Note that I use <code>.narrow(-1, 0, x_size-1)</code> out of habit but <code>[:, :-1]</code> works as well if you prefer that notation.</p>
<pre><code class="lang-python">import torch

A = torch.zeros((3, 3, 3), dtype = torch.float)
X = torch.tensor([[0, 1, 2, 0, 1, 0], [1, 0, 0, 2, 1, 1], [0, 0, 2, 2, 1, 1]])
for a, x in zip(A, X):
    for i, j in zip(x, x[1:]):
        a[i, j] = 1

print(A)



A = torch.zeros((3, 3, 3), dtype = torch.float)
# This code assumes A is contiguous ! If it is not, add
# A = A.contiguous()
# For indexing, collapse the last two dimensions of A
A_view = A.view(A.size(0), -1)
# Compute the indices where you will index in A
x_size = X.size(-1)
indices = X.narrow(-1, 0, x_size-1) * A.stride(1) * A.stride(2) + X.narrow(-1, 1, x_size-1) * A.stride(2)
# Put 1s at the computed indices
A_view.scatter_(1, indices, 1)

print(A)

</code></pre>

Natasha_667: <p>Thanks! works like a charm <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>