🎓 A1Learner
🔍

PyTorch Stacking Tensors

Stacking builds one tensor from several smaller tensors. cat and stack differ in whether the result gains a new dimension.

Join with cat

cat lays the pieces end to end. The result is one longer row.

import torch

first = torch.tensor([20, 22])
second = torch.tensor([19, 21])
print(torch.cat([first, second]))

Stack the same pieces

stack gives each piece its own row. The result is a table.

first = torch.tensor([20, 22])
second = torch.tensor([19, 21])
print(torch.stack([first, second]))

Compare the shapes

The pieces are the same in both calls. Only the shape of the result differs.

first = torch.tensor([20, 22])
second = torch.tensor([19, 21])
print(torch.cat([first, second]).shape)
print(torch.stack([first, second]).shape)

Example

import torch

monday = torch.tensor([21, 23, 19])
tuesday = torch.tensor([20, 22, 24])
wednesday = torch.tensor([18, 20, 21])
week = torch.stack([monday, tuesday, wednesday])
print("Week of readings:")
print(week)
print("Shape:")
print(week.shape)

Key points

  • Use torch.cat(tensors, dim=0) to join tensors along a dimension that already exists
  • Use torch.stack(tensors) to place each tensor along a new first dimension
  • All tensors given to stack must share the same shape
  • cat can join along any existing dimension, rows with dim=0 or columns with dim=1
  • In short, cat makes a dimension longer and stack adds a dimension
  • cat and stack both build a new tensor and leave the pieces unchanged

In this example, monday, tuesday, and wednesday each hold three readings. torch.stack places the three rows into one week tensor. The week prints as a table with one row per day. Its shape is (3, 3), three days of three readings.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Which call adds a new dimension to the result?