Expand the tensor by several dimensions

user_1279459 0
Viewed 2179

In PyTorch, given a tensor of size=[3], how to expand it by several dimensions to the size=[3,2,5,5] such that the added dimensions have the corresponding values from the original tensor. For example, making size=[3] vector=[1,2,3] such that the first tensor of size [2,5,5] has values 1, the second one has all values 2, and the third one all values 3.

In addition, how to expand the vector of size [3,2] to [3,2,5,5]?

One way to do it I can think is by means of creating a vector of the same size with ones-Like and then einsum but I think there should be an easier way.

tensor pytorch
1 Answers

You can first unsqueeze the appropriate number of singleton dimensions, then expand to a view at the target shape with torch.Tensor.expand:

>>> x = torch.rand(3)
>>> target = [3,2,5,5]

>>> x[:, None, None, None].expand(target)

A nice workaround is to use torch.Tensor.reshape or torch.Tensor.view to do perform multiple unsqueezing:

>>> x.view(-1, 1, 1, 1).expand(target)

This allows for a more general approach to handle any arbitrary target shape:

>>> x.view(len(x), *(1,)*(len(target)-1)).expand(target)

For an even more general implementation, where x can be multi-dimensional:

>>> x = torch.rand(3, 2)

# just to make sure the target shape is valid w.r.t to x
>>> assert list(x.shape) == list(target[:x.ndim])

>>> x.view(*x.shape, *(1,)*(len(target)-x.ndim)).expand(target)
user_63313690
Related