How to transform tensor size from [a, b] to [a, b, k] in PyTorch

Viewed 94

I want to add one extra dimension to a tensor and set the value of this dimension as a specific value.
For examples,

print(a.size)
torch.Size([10, 5])
# after tranformation and set a to b
print(b.size)
torch.Size([10, 5, 1])
# after transformation and set a to c
print(c.size)
torch.Size([10, 2, 5])

Thansk in advances.

2 Answers

torch.stack allows you to concatenate two arrays along a new dimension

value = 1.37
a = torch.normal(0, 1, size=(5, 10))
c = torch.stack([a, torch.ones(a.shape) * value], dim=1)
c.shape

Out: torch.Size([5, 2, 10])
c[:, 0, :]

Out: tensor([[-0.2944, -0.7366,  0.6882, -0.7106,  0.0182, -0.1156, -1.0394, -0.7524,
          0.7587, -0.6066],
        [-1.0445, -2.7990,  0.0232,  0.5246, -0.7383,  0.0306, -1.0277, -0.8969,
          0.4026,  0.2006],
        [-1.2622, -0.6563, -1.9218, -0.6932, -1.9633,  1.8271,  0.6753, -0.7564,
          0.0107, -0.2312],
        [-0.8111, -1.0776, -0.8583,  0.2782, -0.8116,  0.0984,  0.4799,  0.6854,
          0.4408, -0.4280],
        [-1.1083,  1.8509,  0.1209,  0.5571, -1.1472,  0.2342,  0.3912,  0.7858,
          0.5879,  0.4139]])

c[:, 1, :]

Out: tensor([[1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700,
         1.3700],
        [1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700,
         1.3700],
        [1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700,
         1.3700],
        [1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700,
         1.3700],
        [1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700, 1.3700,
         1.3700]])

First part of the question can be solved with unsqueeze

a.shape
>>>torch.Size([10, 5])

a.unsqueeze(2).shape
>>>torch.Size([10, 5, 1]))

You can also use a.unsqueeze(-1) to add dimension on the last dimension, but better to use 2, as in The Zen of Python: Explicit is better than implicit.

Related