unbind and stack

torch.unbind removes the designated dimension from a tensor and returns a tuple of lower-dimensional slices along that axis. torch.stack does the opposite. It joins a sequence tensors along a new dimension.

unbind

Visually unbind slices a tensor along a dimension. For example, for 2-d tensors, slicing along dim 0 is breaking the tensor into rows, while slicing along dim 1 is breaking the tensor into cols. However “slicing” is hard to picture in higher dimensions. What exactly is the logic of the slicing action?

First, let’s see an example of unbind for a 2-d tensor.

import torch
 
x = torch.tensor([
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
])
print(x.shape)

torch.Size([3, 3])

x_unbind_0 = x.unbind(dim = 0)
print(x_unbind_0, len(x_unbind_0), x_unbind_0[0].shape)
x_unbind_1 = x.unbind(dim = 1)
print(x_unbind_1, len(x_unbind_1), x_unbind_1[0].shape)

(tensor([1, 2, 3]), tensor([4, 5, 6]), tensor([7, 8, 9])) 3 torch.Size([3])
(tensor([1, 4, 7]), tensor([2, 5, 8]), tensor([3, 6, 9])) 3 torch.Size([3])

Dimension wise, unbind reduces the rank of the tensor by 1, so a 2-d tensor becomes an array of 1-d tensors. This holds true no matter which dimension you unbind. The dimension parameter only changes the shape, not the final rank of the result. To see what slicing actually is, let’s use a more complex example, a tensor with dimension 4.

tensor = torch.randint(low=1, high=11, size=(2, 1, 3, 4))
tensor

tensor([[[[ 1,  1,  7,  3],
          [ 7,  6,  5,  3],
          [10,  8,  7,  7]]],


        [[[ 6,  3,  3, 10],
          [10,  2,  2,  8],
          [ 7,  7,  9,  8]]]])

Let’s say we want to unbind dimension 2, which is the 3rd dimension. How do we visualize the slicing action? We can do so by looking at a single element (say index = 0) in that dimension. We can achieve this by using the bracket notation (also called “slicing” ) or a function called select, which is a simplified version of the bracket notation. The doc describes select as “slices the input tensor along the selected dimension at the given index”. Either way, you will see an element of the result array after the unbind. Instead of giving you the whole array, this allows you to see a single element.

tensor[:, :, 0, :]

tensor([[[ 1,  1,  7,  3]],

        [[ 6,  3,  3, 10]]])

tensor.select(2, 0)

tensor([[[ 1,  1,  7,  3]],

        [[ 6,  3,  3, 10]]])

As you can see below, at index 0, the element of the result array is the same as the select function.

tensor.unbind(dim=2)

(tensor([[[ 1,  1,  7,  3]],
 
         [[ 6,  3,  3, 10]]]),
 tensor([[[ 7,  6,  5,  3]],
 
         [[10,  2,  2,  8]]]),
 tensor([[[10,  8,  7,  7]],
 
         [[ 7,  7,  9,  8]]]))

stack

Stack is the opposite of the unbind, where you glue togehter a sequence of tensor of the same size along a dimension. Here the dimensionality of the operation has more variations. Here is an example of stacking two tensors along different dimensions. Here both t1 and t2 have 3 dimensions. You can stack them on any of these dimensions.

t1 = torch.randint(low=1, high=11, size=(2, 3))
t2 = torch.randint(low=1, high=11, size=(2, 3))
print(t1, '\n', t2)

tensor([[3, 1, 8],
        [5, 2, 3]]) 
 tensor([[5, 7, 6],
        [4, 9, 9]])

res_0 = torch.stack([t1, t2], dim = 0)
print(res_0, '\n', res_0.shape)

tensor([[[3, 1, 8],
         [5, 2, 3]],

        [[5, 7, 6],
         [4, 9, 9]]]) 
 torch.Size([2, 2, 3])

res_1 = torch.stack([t1, t2], dim = 1)
print(res_1, '\n', res_1.shape)

tensor([[[3, 1, 8],
         [5, 7, 6]],

        [[5, 2, 3],
         [4, 9, 9]]]) 
 torch.Size([2, 2, 3])

res_2 = torch.stack([t1, t2], dim = 2)
print(res_2, '\n', res_2.shape)

tensor([[[3, 5],
         [1, 7],
         [8, 6]],

        [[5, 4],
         [2, 9],
         [3, 9]]]) 
 torch.Size([2, 3, 2])

Under the hood, stack does two things, first, for each of the element, it creates a new dimension at the designation dimension (with unsqueez). Then it concat the elements along that dimension (with cat). So dim = 1 becomes a new dimension, dim = 0 and dim = 2are the original 2-d dimensions.

torch.cat([t1.unsqueeze(1), t2.unsqueeze(1)], dim=1)

tensor([[[3, 1, 8],
         [5, 7, 6]],

        [[5, 2, 3],
         [4, 9, 9]]])

Think of stack glueing together several tensors in an extra dimension. From the new dimension, each index holds a original tensor, but from the old dimension, you now see you can see how they are glued together. In the following example. res at dim=0 shows you the original t1 and t2 at index 0 and 1.

res = torch.stack([t1, t2], dim = 0)
torch.select(res, dim=0, index=0)

tensor([[3, 1, 8],
        [5, 2, 3]])

At the original dimension, dim = 1, which is the original dim 0 of the 2-d dimension, the index = 0 still denotes “first row in the matrix”, so you get 2 first rows in t1 and t2 togehter.

torch.select(res, dim=1, index=0)

tensor([[3, 1, 8],
        [5, 7, 6]])

Finally, see that unbind then stack produces the original tensor

seq = tensor.unbind(dim=2)
res_4d = torch.stack(seq, dim = 2)
print(res_4d, '\n', res_4d.shape, '\n', torch.equal(tensor, res_4d))

tensor([[[[ 1,  1,  7,  3],
          [ 7,  6,  5,  3],
          [10,  8,  7,  7]]],


        [[[ 6,  3,  3, 10],
          [10,  2,  2,  8],
          [ 7,  7,  9,  8]]]]) 
 torch.Size([2, 1, 3, 4]) 
 True