Left shift tensor in PyTorch

user_6997665 0
Viewed 3272

I have a tensor a of shape (1, N, 1). I need to left shift the tensor along dimension 1 and add a new value as replacement. I have found a way to make this work and following is the code.

a = torch.from_numpy(np.array([1, 2, 3]))
a = a.unsqueeze(0).unsqeeze(2)  # (1, 3, 1), my data resembles this shape, therefore the two unsqueeze
# want to left shift a along dim 1 and insert a new value at the end
# I achieve the required shifts using the following code
b = a.squeeze
c = b.roll(shifts=-1)
c[-1] = 4
c = c.unsqueeze(0).unsqueeze(2)
# c = [[[2], [3], [4]]]

My question is, is there a simpler way to do this? Thanks.

torch python-3.x
1 Answers

You don't actually need to squeeze and perform your operations followed by unsqeezing your input tensor a. Instead you can directly do those two operations as follows:

# No need to squeeze
c = torch.roll(a, shifts=-1, dims=1)
c[:,-1,:] = 4
# No need to unsqeeze
# c = [[[2], [3], [4]]]
user_51365500
Related