PyTorch torch.cat vs torch.stack: shapes and examples

Use torch.cat to join tensors along an existing dimension. Use torch.stack to insert a new dimension that separates the inputs. For two matrices of shape (2, 3), cat can produce (4, 3) or (2, 6); stack produces a three-dimensional tensor.

The important question is what each axis represents: samples, time steps, channels, features, or something else. A valid output shape does not prove that the values are arranged as the next layer expects.

This guide follows five original examples with captured output. We ran 36 checks using Python 3.12.14 and PyTorch 2.15.0+cpu. Integer fixtures make the arrangement visible; float64 fixtures verify gradients. These are CPU behavior checks, not performance measurements.

Join rows or columns with torch.cat

Both inputs below have two rows and three columns. The values in the second matrix are larger so it is easy to see where each source lands.

import torch

a = torch.tensor([[1, 2, 3], [4, 5, 6]])
b = torch.tensor([[10, 20, 30], [40, 50, 60]])
rows = torch.cat((a, b), dim=0)
columns = torch.cat((a, b), dim=1)
last_axis = torch.cat((a, b), dim=-1)
print("dim=0:", tuple(rows.shape), rows.tolist())
print("dim=1:", tuple(columns.shape), columns.tolist())
print("dim=-1 equals dim=1:", torch.equal(last_axis, columns))

dim=0: (4, 3) [[1, 2, 3], [4, 5, 6], [10, 20, 30], [40, 50, 60]]
dim=1: (2, 6) [[1, 2, 3, 10, 20, 30], [4, 5, 6, 40, 50, 60]]
dim=-1 equals dim=1: True

With dim=0, the row counts add: two plus two gives four, while the column count stays three. With dim=1, the column counts add, producing two rows of six values. The operation keeps two dimensions in both cases.

For this two-dimensional input, dim=-1 means the last axis and matches dim=1. It does not always mean columns in a larger tensor; it means whichever axis is last in that tensor’s shape.

The torch.cat reference requires the non-concatenated dimensions to match, with a documented exception for a one-dimensional empty tensor of shape (0,). The next examples show the practical boundary of that rule.

Create a new axis with torch.stack

Stacking preserves a separate position for each input along the inserted axis. All input shapes must match. In the torch.stack API, dim selects where that new axis goes.

import torch

a = torch.tensor([[1, 2, 3], [4, 5, 6]])
b = torch.tensor([[10, 20, 30], [40, 50, 60]])
front = torch.stack((a, b), dim=0)
middle = torch.stack((a, b), dim=1)
back = torch.stack((a, b), dim=-1)
via_cat = torch.cat((a.unsqueeze(1), b.unsqueeze(1)), dim=1)
print("dim=0:", tuple(front.shape), front.tolist())
print("dim=1:", tuple(middle.shape), middle.tolist())
print("dim=-1:", tuple(back.shape), back.tolist())
print("unsqueeze then cat equals stack:", torch.equal(via_cat, middle))

dim=0: (2, 2, 3) [[[1, 2, 3], [4, 5, 6]], [[10, 20, 30], [40, 50, 60]]]
dim=1: (2, 2, 3) [[[1, 2, 3], [10, 20, 30]], [[4, 5, 6], [40, 50, 60]]]
dim=-1: (2, 3, 2) [[[1, 10], [2, 20], [3, 30]], [[4, 40], [5, 50], [6, 60]]]
unsqueeze then cat equals stack: True

At dim=0, the first output slice is the whole first matrix, and the second slice is the whole second matrix. At dim=1, each original row contains a pair of source rows. Both outputs happen to have shape (2, 2, 3), because the number of inputs and the original row count are both two. Their values are arranged differently.

At dim=-1, the output has shape (2, 3, 2). Each final pair contains corresponding values from the two inputs, such as [1, 10]. For these matrices, stack can insert at positive axes 0, 1, or 2. Cat has only the existing axes 0 and 1 available.

The final line verifies one useful identity: inserting a singleton axis into each input, then concatenating along it, reproduces this stack operation. unsqueeze adds an axis of length one and shares the source tensor’s underlying data. Adding that axis does not rearrange the original values.

Choose the axis from the meaning of the data

If two tensors contain batches of feature vectors shaped (batch, features), cat at dim=0 combines the batches while preserving feature width. Cat at dim=1 joins features for corresponding rows. That second use requires the rows to describe the same samples in the same order.

Shape checks cannot establish that correspondence. Two batches can have equal lengths while containing unrelated samples. Joining their feature columns would still run, but the resulting records would have the wrong meaning.

If each input is one equally shaped sample, stack at dim=0 can create a batch axis. If the inputs are already batches, the same call adds an extra axis around those batches. Inspect the tensors returned by your PyTorch DataLoader before deciding which case you have.

Handle unequal lengths without assuming broadcasting

A shorter final batch can still be concatenated when the feature width matches:

import torch

a = torch.tensor([[1, 2, 3], [4, 5, 6]])
short = torch.tensor([[7, 8, 9]])
joined = torch.cat((a, short), dim=0)
print("different batch sizes:", tuple(joined.shape), joined.tolist())
errors = {}
for label, operation in [
    ("stack unequal shapes", lambda: torch.stack((a, short))),
    ("cat unequal non-join axes", lambda: torch.cat((a, short), dim=1)),
    ("cat missing axis", lambda: torch.cat((a, a), dim=2)),
]:
    try:
        operation()
    except (RuntimeError, IndexError) as exc:
        errors[label] = type(exc).__name__
        print(label + ":", type(exc).__name__)

different batch sizes: (3, 3) [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
stack unequal shapes: RuntimeError
cat unequal non-join axes: RuntimeError
cat missing axis: IndexError

The first call combines a (2, 3) tensor and a (1, 3) tensor into (3, 3). Stacking those tensors fails because their shapes differ. Concatenating them at dim=1 also fails: the non-concatenated row axis has lengths two and one.

Cat does not repeat the one-row input to match the other tensor. The error is useful evidence that the intended alignment needs attention. Padding, packing, or explicitly repeating values would be separate choices, each with consequences for the data.

The last call tries to concatenate at dim=2, but these inputs have only two axes. It raises IndexError. Error wording can vary by version; this example prints the exception class from the tested runtime.

Our additional checks rejected an empty list passed to either function. A tensor with zero rows is a different case: a (0, 3) tensor concatenated with the (2, 3) fixture at dim=0 added no rows. We also verified the documented (0,) empty-tensor exception for cat. Keep an empty collection of inputs separate from an input tensor that contains no elements.

Stack scalars, or add an axis before concatenating

A scalar tensor has shape (), with no existing axis to extend. The following example makes that difference explicit:

import torch

x = torch.tensor(3.0)
y = torch.tensor(7.0)
stacked = torch.stack((x, y))
promoted = torch.cat((x.unsqueeze(0), y.unsqueeze(0)))
print("scalar input shape:", tuple(x.shape))
print("stacked:", tuple(stacked.shape), stacked.tolist())
print("promoted cat:", tuple(promoted.shape), promoted.tolist())
try:
    torch.cat((x, y))
except RuntimeError as exc:
    scalar_error = type(exc).__name__
    print("cat scalars:", scalar_error)

scalar input shape: ()
stacked: (2,) [3.0, 7.0]
promoted cat: (2,) [3.0, 7.0]
cat scalars: RuntimeError

Stack creates a new vector axis and returns [3.0, 7.0]. Direct cat fails. Adding a singleton axis to each scalar first gives two one-element vectors, which cat can combine into the same result.

This is useful when collecting tensor-valued scalar results. Decide whether you need to retain their gradient connections before converting values to ordinary Python numbers; the gradient example below keeps tensors throughout.

Check how gradients return to the inputs

The joined result can participate in an autograd calculation. We use two different losses so that the expected input gradients can be checked independently.

import torch

x = torch.tensor([1.0, 2.0], dtype=torch.float64, requires_grad=True)
y = torch.tensor([3.0, 4.0], dtype=torch.float64, requires_grad=True)
torch.cat((x, y)).square().sum().backward()
cat_x_grad = x.grad.clone()
cat_y_grad = y.grad.clone()
print("cat gradients:", cat_x_grad.tolist(), cat_y_grad.tolist())

x.grad = None
y.grad = None
weights = torch.tensor([[1.0, 10.0], [100.0, 1000.0]], dtype=torch.float64)
(torch.stack((x, y), dim=1) * weights).sum().backward()
stack_x_grad = x.grad.clone()
stack_y_grad = y.grad.clone()
print("stack gradients:", stack_x_grad.tolist(), stack_y_grad.tolist())

x.grad = None
torch.cat((x, x)).sum().backward()
print("same input used twice:", x.grad.tolist())

cat gradients: [2.0, 4.0] [6.0, 8.0]
stack gradients: [1.0, 100.0] [10.0, 1000.0]
same input used twice: [2.0, 2.0]

For cat, the loss is the sum of squared joined values. Its derivative is twice each value, so x receives [2, 4] and y receives [6, 8]. For stack at dim=1, x occupies the first column and y the second. Multiplying by the weight matrix gives x gradients [1, 100] and y gradients [10, 1000].

We clear the old gradients between calculations. In the final calculation, x appears twice in the concatenation, so each element receives two contributions of one. The resulting gradient is [2, 2]. The tests also confirm that backward leaves the source values unchanged.

Joining tensors is separate from accumulating gradients across microbatches. Our gradient-accumulation guide covers loss weighting and optimizer-step boundaries. Combining tensor values alone does not implement that training loop.

Verify one known value as well as the shape

Before joining real data, write down the meaning of each axis and the expected output shape. Then trace one known value from each input to its destination. The two stack outputs with shape (2, 2, 3) show why checking shape alone can miss an axis mistake.

These tests cover dense CPU tensors, selected invalid inputs, and exact gradient routing. A noncontiguous input also worked in a separate cat check. We did not test CUDA performance, sparse or quantized tensors, mixed-device inputs, or autograd with the out argument. Keep performance and training-quality claims separate from these value and shape checks.