In NumPy, I would do
a = np.zeros((4, 5, 6))
a = a[:, :, np.newaxis, :]
assert a.shape == (4, 5, 1, 6)
How to do the same in PyTorch?
Add a comment
|
3 Answers
a = torch.zeros(4, 5, 6)
a = a[:, :, None, :]
assert a.shape == (4, 5, 1, 6)
-
11And for what it's worth
np.newaxisis justNone, anyway. Commented Dec 27, 2020 at 21:53 -
1
-
2In situations where you don't have to do any indexing on other axes when requiring this operation, I don't see this being very useful. In my opinion
a.unsqueeze(2)is much more effective and to the point.– IvanCommented Dec 27, 2020 at 22:54 -
@Ivan I see how it can be more readable to some. Why is it more
effective?– GulzarCommented Dec 28, 2020 at 9:42 -
I didn't mean in terms of speed and performance of course. What I meant was it's a bit troublesome if you have a lot of dimensions and are not looking to do any slicing on other dims at the same time you're adding that new dim. But, we can agree it does the exact same thing ;)– IvanCommented Dec 28, 2020 at 9:49
You can add a new axis with torch.unsqueeze() (first argument being the index of the new axis):
>>> a = torch.zeros(4, 5, 6)
>>> a = a.unsqueeze(2)
>>> a.shape
torch.Size([4, 5, 1, 6])
Or using the in-place version: torch.unsqueeze_():
>>> a = torch.zeros(4, 5, 6)
>>> a.unsqueeze_(2)
>>> a.shape
torch.Size([4, 5, 1, 6])
x = torch.tensor([1, 2, 3, 4])
y = torch.unsqueeze(x, 0)
y will be -> tensor([[ 1, 2, 3, 4]])
EDIT: see more details here: https://pytorch.org/docs/stable/generated/torch.unsqueeze.html
answered Dec 27, 2020 at 21:55
