cat, stack

cat glues tensors along an axis that already exists; stack creates a new axis and glues along it. Choosing correctly is the difference between a (6, 4) and a (2, 3, 4), and between a real batch and a shape that only looks plausible.

The idea

torch.cat(tensors, dim=0) concatenates along an existing dimension. Every other dimension must match exactly. Two (2, 3) tensors cat along dim=0 into (4, 3); along dim=1 into (2, 6).

torch.stack(tensors, dim=0) inserts a new dimension whose size is the number of tensors. Two (3, 4) tensors stack at dim=0 into (2, 3, 4), or at dim=1 into (3, 2, 4).

• cat requires matching shapes except along dim.
• stack requires all tensors to share the exact same shape.
• torch.stack([a, b]) equals torch.cat([a.unsqueeze(0), b.unsqueeze(0)]).

Shapes tell you which to reach for. A list of per-sample outputs belongs in stack; pooling feature channels or merging batches belongs in cat. The inverses are torch.split and torch.chunk, which cut one tensor along a dimension.

Worked example

Two attention heads each produce (8, 16) — 8 tokens, width 16.

• torch.cat([h1, h2], dim=1) → (8, 32): a pure feature concatenation.
• torch.stack([h1, h2], dim=0) → (2, 8, 16): heads kept as a separate axis.
• .transpose(0, 1).reshape(8, 32) on that stack → (8, 32), the same result as cat, but only after moving the head axis next to the feature axis.

Now batch three images of shape (3, 224, 224). torch.stack(imgs) → (3, 3, 224, 224), a proper batch. torch.cat(imgs, dim=0) → (9, 224, 224), which treats each image's rows as separate entries. It runs without error and is almost always wrong — a reliable signal you wanted stack.

In code

import torch

a = torch.zeros(2, 3)
b = torch.ones(4, 3)
print(torch.cat([a, b], dim=0).shape)     # torch.Size([6, 3])

c = torch.zeros(2, 3)
print(torch.stack([a, c], dim=0).shape)   # torch.Size([2, 2, 3])
print(torch.stack([a, c], dim=1).shape)   # torch.Size([2, 2, 3])

h1 = torch.randn(8, 16)
h2 = torch.randn(8, 16)
print(torch.cat([h1, h2], dim=1).shape)   # torch.Size([8, 32])
heads = torch.stack([h1, h2], dim=0)      # (2, 8, 16)
print(heads.transpose(0, 1).reshape(8, 32).shape)   # torch.Size([8, 32])

Check yourself

1. You have 10 tensors of shape (3, 32, 32). Which op builds (10, 3, 32, 32), and which builds (3, 320, 32)?
2. What error does stack raise on tensors with different shapes, and which op accepts them instead?
3. If a and b are both (2, 3), what shapes do cat and stack produce at dim=0?

Key takeaways

• cat extends an existing axis; stack adds a new one.
• cat allows differing sizes along dim; stack requires identical shapes.
• To get a batch from a list of samples, use stack.