Adding an axis, or extending one

three tensors of shape (4,)

torch.stack(xs)        ->  (3, 4)     a new axis
torch.cat(xs)          ->  (12,)      the existing axis, extended

stack requires every input to have the same shape and adds an axis of length len(xs). cat requires them to match on every axis except the one being joined, and extends that one.

Batching individual examples is stack. Joining batches, or the outputs of several heads along the feature axis, is cat.

The shape that hides it

With three tensors of shape (4,), stack gives (3, 4) and cat gives (12,), which are obviously different. With three of shape (1, 4):

torch.stack(xs)   ->  (3, 1, 4)
torch.cat(xs)     ->  (3, 4)

and (3, 4) is what you wanted. So a cat that should have been a stack works perfectly as long as every example arrives with a leading axis of 1, and produces (3, 1, 4) versus (3, 4) the moment one does not. Data loaders are full of exactly this.

Neither is free

Both allocate. stack is unsqueeze on each input followed by cat, so it is a copy of everything either way, and there is no view-based version: separate tensors live in separate buffers, and one tensor needs one buffer.

That is worth knowing rather than worrying about. Collating a batch is a real copy and always was; the thing to avoid is doing it more than once per batch.

The related pair

torch.chunk(x, n, dim)      # split into n pieces, views
torch.split(x, size, dim)   # split into pieces of a given size, views

Both return views, since splitting along an axis is a stride and offset change. Joining allocates and splitting does not, which is the asymmetry to remember.

Your task

def collate(examples: list) -> torch.Tensor

Each example is a 1-D tensor of the same length. Return a batch of shape (len(examples), length).

The starter uses cat.