Udemy - Section 2 - Lecture 32: Differences between torch.stack, torch.vstack and torch.hstack

[LUS24]

Hi, I just wanted to mention that in the lecture video in Section 2 - Lecture 32, timestamp 3:28, the following is said:

Okay, so maybe not stack or V stack.
We can just define what dimension we'd like to combine them on.
I wonder if there is a torch v stack towards V stack.
Oh, there is.
And is there a torch stack for horizontal stack?
Here is a hstack.
Beautiful.
So we'll focus on just the plane stack.
If you want to have a look at VH Stack, it'll be quite similar to what we're going to do with Stack.
And same with Stack.

This made me think that setting dim=0 or dim=1 in toch.stack would be equivalent to torch.vstack or torch.hstack.

After making a few tests I saw that this is not the case. toch.stack creates a new dimension and torch.vstack and torch.hstack don't, here are a few examples:

# Starting tensors
tensor_A = torch.tensor(
    [
        [1, 2],
        [3, 4]
    ]
)

tensor_B = torch.tensor(
    [
        [5, 6],
        [7, 8]
    ]
)

# torch.stack dim=0
tensor_stack = torch.stack((tensor_A, tensor_B), dim=0)
print("Shapes: ", tensor_A.shape, tensor_B.shape, tensor_stack.shape)
print("Dimensions: ", tensor_A.ndim, tensor_B.ndim, tensor_stack.ndim)
tensor_stack

Shapes:  torch.Size([2, 2]) torch.Size([2, 2]) torch.Size([2, 2, 2])
Dimensions:  2 2 3

tensor([[[1, 2],
         [3, 4]],

        [[5, 6],
         [7, 8]]])

# torch.stack dim=1
tensor_stack = torch.stack((tensor_A, tensor_B), dim=1)
print("Shapes: ", tensor_A.shape, tensor_B.shape, tensor_stack.shape)
print("Dimensions: ", tensor_A.ndim, tensor_B.ndim, tensor_stack.ndim)
tensor_stack

Shapes:  torch.Size([2, 2]) torch.Size([2, 2]) torch.Size([2, 2, 2])
Dimensions:  2 2 3

tensor([[[1, 2],
         [5, 6]],

        [[3, 4],
         [7, 8]]])

# torch.vstack
tensor_stack = torch.vstack((tensor_A, tensor_B))
print("Shapes: ", tensor_A.shape, tensor_B.shape, tensor_stack.shape)
print("Dimensions: ", tensor_A.ndim, tensor_B.ndim, tensor_stack.ndim)
tensor_stack

Shapes:  torch.Size([2, 2]) torch.Size([2, 2]) torch.Size([4, 2])
Dimensions:  2 2 2

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

# torch.hstack
tensor_stack = torch.hstack((tensor_A, tensor_B))
print("Shapes: ", tensor_A.shape, tensor_B.shape, tensor_stack.shape)
print("Dimensions: ", tensor_A.ndim, tensor_B.ndim, tensor_stack.ndim)
tensor_stack

Shapes:  torch.Size([2, 2]) torch.Size([2, 2]) torch.Size([2, 4])
Dimensions:  2 2 2

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

From the results it seems that stack creates a new dimension resulting in a 3D tensor while vstack and hstack preserve the 2 dimensions but increase the 2D tensor shape from (2,2) to either (4,2) or (2,4).

The equivalent to vstack and hstack would be torch.cat(tensor, dim=0) or torch.cat(tensor, dim=1)

Maybe a quick note could be added to the lecture.

Thanks for the great content!