I have a tensor A of shape [4, 4096] and another tensor B of shape [5, 4096].

1. I want to take each element of B along the 0th axis, duplicate it and stack it along the tensor A's first and last index along 0th axis (The resulting tensor after this operation should be the size of [6, 4096]).

2. I want to repeat this operation for each tensor in tensor B. The final tensor should be the size of [30, 4096].

A way to visualize it...

A = [A1, A2, A3, A4]
B = [B1, B2, B3, B4, B5]

T = [ B1, A1, A2, A3, A4, B1, B2, A1, A2, A3, A4, B2, B3, A1, A2, A3, A4, B3,.... B5, A1, A2, A3, A4, B5 ]

Here,
B = [5, 4096] so, B1 = [4096]
A = [4, 4096] so, A1 = [4096]
T = [30, 4096]

CodePudding user response：

Well, I did find an answer for this.

A = torch.rand(4, 4096)
B = torch.rand(5, 4096)

tens_arr = []

for seq in B:
 
    #creates a tensor of [6, 4096]
    t = torch.cat((seq.unsqueeze(dim=0), A, seq.unsqueeze(dim=0)))

    tens_arr.append(t) 

 positional_tensor = torch.cat(T, dim=0)

 print(positional_tensor.shape) #Prints [30, 4096]

Is there any way for me to further make this code shorter or even not use the LOOP there at all?

CodePudding user response：

You can take advantage of tensor views to do this efficiently (create a list of views, and run one actual concatenation op at the end).

import torch
N = 4096
A, B = torch.zeros(4, N), torch.ones(5, N)
T = torch.cat([c for b in B.unsqueeze(1) for c in (b, A, b)], 0)
print(T.shape, T[:, 0])

This prints the correct result

torch.Size([30, 4096]) tensor([1., 0., 0., 0., 0., 1., 1., 0., 0., 0., 0., 1., 1., 0., 0., 0., 0., 1.,
        1., 0., 0., 0., 0., 1., 1., 0., 0., 0., 0., 1.])

and, at least on my machine, runs in ~half the time of the version which concatenates each length-6 tensor separately.