Iterate through tensor dimensions

username: ptrch_c_m
cooked: <p>I have a 3D tensor and I want to iterate through the second dimension. The easiest approach is to use a for loop on the second dimension. My question is what is the most efficient/appropriate way to do it. For example:</p>
<pre><code class="lang-auto">t = torch.rand(2, 5, 3)

print(t)
tensor([[[0.1582, 0.7643, 0.1352],
         [0.7157, 0.4503, 0.7935],
         [0.9954, 0.6714, 0.0386],
         [0.7655, 0.2867, 0.0764],
         [0.2259, 0.3079, 0.8390]],

        [[0.2655, 0.2578, 0.8421],
         [0.7871, 0.7559, 0.2978],
         [0.4458, 0.6559, 0.9398],
         [0.6346, 0.4585, 0.3734],
         [0.8566, 0.1053, 0.0559]]])

for idx in range(t.shape[1]):
    print(t[:,idx,:])
tensor([[0.1582, 0.7643, 0.1352],
        [0.2655, 0.2578, 0.8421]])
tensor([[0.7157, 0.4503, 0.7935],
        [0.7871, 0.7559, 0.2978]])
tensor([[0.9954, 0.6714, 0.0386],
        [0.4458, 0.6559, 0.9398]])
tensor([[0.7655, 0.2867, 0.0764],
        [0.6346, 0.4585, 0.3734]])
tensor([[0.2259, 0.3079, 0.8390],
        [0.8566, 0.1053, 0.0559]])
</code></pre>

username: ptrblck
cooked: <p>If you want to iterate through a tensor dimension, using a for loop is a valid approach. However, it’s unclear if the loop is an optimal approach since you did not explain what your actual goal is besides iterating.</p>

username: ptrch_c_m
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> thank you for your answer. Let my clarify what the tensor represents. The tensor is a batch of sequences, where each sequence element is features, thus the tensor is</p>
<p><code>[batch_size, sequence_length, features]</code></p>
<p>What I want, is to process each sequence element separately in batches, for example to pass them through a linear layer etc. That is why I want the for loop. I wonder if there is a better way of doing this.</p>

username: ptrblck
cooked: <p><code>nn.Linear</code> layers accept inputs tensors with additional dimensions and will loop over these internally, as seen here:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn

# setup
batch_size, sequence_length, features = 2, 3, 4
x = torch.randn(batch_size, sequence_length, features)

lin = nn.Linear(features, features*2)

# loop approach
outputs = []
for i in range(sequence_length):
    x_ = x[:, i]
    out = lin(x_)
    outputs.append(out)
outputs = torch.stack(outputs, dim=1)
print(outputs.shape)
# torch.Size([2, 3, 8])

# single call
ref = lin(x)
print(ref.shape)
# torch.Size([2, 3, 8])

# compare
print((outputs - ref).abs().max())
# tensor(0., grad_fn=&lt;MaxBackward1&gt;)
</code></pre>
<p>so if you are using <code>nn.Linear</code> layers you would not have to iterate the input tensor manually.</p>

username: ptrch_c_m
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> thank you once again for your detailed answer.</p>
<p>I am building something more complex, a version of an RNN model, specifically, a VRNN model (<a href="https://arxiv.org/pdf/1506.02216" rel="noopener nofollow ugc">arxiv</a>), where the computation of earlier sequence elements are used by the computations of the following elements, thus they cannot happen simultaneously.</p>
<p>Generally, I think that a for loop is the most basic way of doing things, which is the reason of the post.I guess that, in this case, this for loop is the best you can do.</p>