Adding an axis, or extending one

stack adds a new axis. cat joins values along an existing axis. For a list of equally sized vectors, that distinction determines whether you get a batch or one long vector:

import torch

examples = [torch.tensor([1., 2.]), torch.tensor([3., 4.])]
batch = torch.stack(examples)
joined = torch.cat(examples)
print(batch.tolist())  # [[1.0, 2.0], [3.0, 4.0]]
print(list(batch.shape))  # [2, 2]
print(joined.tolist())  # [1.0, 2.0, 3.0, 4.0]
print(list(joined.shape))  # [4]

stack requires all input shapes to match. Its new axis has one position per input tensor. For ordinary nonempty inputs, cat requires matching sizes on every axis except the one being joined.

Does the input already have a batch axis?

For three individual examples of shape (4,), stacking produces (3, 4). For three single-example batches of shape (1, 4), concatenating along axis zero produces (3, 4) instead:

Choose based on what each axis represents. Stacking is appropriate for individual examples without a batch axis. Concatenation is appropriate for joining existing batches or extending a feature axis.

Storage and copying

With the default calls used here, both operations construct a new result buffer and copy the input values into it. Stacking has the same value and shape semantics as inserting an axis into each input and concatenating; that is not a guarantee about its internal implementation.

The harness creates each example separately, so the batch needs combined storage. In other code, inputs might already be slices of a shared batch; retaining that original batch can avoid an unnecessary split-and-rejoin. Separate tensor objects do not necessarily mean separate storage.

Splitting a strided tensor with split or chunk returns views. split can use a piece size or explicit section sizes. chunk requests a number of pieces, but may return fewer than requested. Neither requires copying the values merely to describe the pieces.

Give each tensor axis a meaning

A tensor's shape tells you how many entries each axis has. It does not tell you what those entries represent. For example, (batch, features) and (features, batch) use different axis orders. Name the axes before choosing an operation.

A reduction combines entries along an axis, for example by summing or averaging. Averaging the features of a (B, D) tensor gives one number per example, with shape (B,). Keeping the reduced axis gives shape (B, 1) instead.

For [[2, 4, 6], [1, 3, 5]], the row means with the axis kept are [[4], [3]]. Subtracting them from the input gives [[-2, 0, 2], [-2, 0, 2]]: each row has been centered using its own mean.

Reshape and transpose do different jobs

A reshape regroups elements using the operation's traversal order. A transpose changes the axis order. For the usual row-major reshape, [[1, 2, 3], [4, 5, 6]] reshaped to (3, 2) becomes [[1, 2], [3, 4], [5, 6]]. Transposing the original instead gives [[1, 4], [2, 5], [3, 6]].

Both results have shape (3, 2), but their entries mean different things. Check values as well as shapes. Unequal axis sizes make many axis mistakes easier to spot than square examples do.

Concatenate along an axis or stack on a new one

Concatenation extends an existing axis. Stacking creates a new axis. For two arrays of shape (2, 3), concatenating on axis zero gives (4, 3). Stacking on a new leading axis gives (2, 2, 3).