Stacking tensors from multiple tensors in a specified indexing order

mks1596: <p>Hi,</p>
<p>I want to stack two tensors along a dimension, but not sequentially. Rather I want to specify the particular indexing of the stacking along that dimension. As a concrete example, I will show how this should work for just two tensors A and B:</p>
<pre><code class="lang-auto">A=torch.tensor([[1,2,3],
                [4,5,6],
                [7,8,9]])
B=torch.tensor([[10,11,12],
                [13,14,15],
                [16,17,18],
                [19,20,21]])
</code></pre>
<p>Stacking A and B along dim = 0 with this indexing for interleaving A with B:</p>
<pre><code class="lang-auto">stack_indexing_A = torch.tensor([1,3,6])
</code></pre>
<p>should yield this output:</p>
<pre><code class="lang-auto">result = torch.tensor([[10,11,12], 
                       [1,2,3],
                       [13,14,15],
                       [4,5,6],
                       [16,17,18],
                       [19,20,21],
                       [7,8,9]])
</code></pre>
<p>Basically I interleave A with B in a specified index ordering for one of the two (A in this case). This needs to be done without allocating a torch.empty tensor and filling it out because both these tensors are part of the computational graph of my model and assigining a new torch empty tensor might cause it to be rendered out of the computation graph. I am trying to do this with torch gather but as far as I understand that only works with a single tensor. Thanks for any help!</p>

ptrblck: <aside class="quote no-group" data-username="mks1596" data-post="1" data-topic="199026">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/9de053/48.png" class="avatar"> mks1596:</div>
<blockquote>
<p>This needs to be done without allocating a torch.empty tensor and filling it out because both these tensors are part of the computational graph of my model and assigining a new torch empty tensor might cause it to be rendered out of the computation graph.</p>
</blockquote>
</aside>
<p>Why would this be the case?<br>
It works fine in my minimal example:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch

A = torch.tensor([[1,2,3],
                  [4,5,6],
                  [7,8,9]]).float().requires_grad_()
B = torch.tensor([[10,11,12],
                  [13,14,15],
                  [16,17,18],
                  [19,20,21]]).float().requires_grad_()

stack_indexing_A = torch.tensor([1,3,6])
stack_mask_A = torch.tensor([False]*(A.size(0)+B.size(0)))
stack_mask_A[stack_indexing_A] = True

C = torch.empty(A.size(0)+B.size(0), A.size(1))
C[stack_mask_A] = A
C[~stack_mask_A] = B

print(C)
# tensor([[10., 11., 12.],
#         [ 1.,  2.,  3.],
#         [13., 14., 15.],
#         [ 4.,  5.,  6.],
#         [16., 17., 18.],
#         [19., 20., 21.],
#         [ 7.,  8.,  9.]], grad_fn=&lt;IndexPutBackward0&gt;)

C.mean().backward()
print(A.grad)
# tensor([[0.0476, 0.0476, 0.0476],
#         [0.0476, 0.0476, 0.0476],
#         [0.0476, 0.0476, 0.0476]])
print(B.grad)
# tensor([[0.0476, 0.0476, 0.0476],
#         [0.0476, 0.0476, 0.0476],
#         [0.0476, 0.0476, 0.0476],
#         [0.0476, 0.0476, 0.0476]])
</code></pre>