PyTorch Tensor Operations

PyTorch tensor operations are the vocabulary you use to write models. Nearly everything a neural network does is a combination of element-wise math, matrix multiplication and reshaping. Getting comfortable with shapes is the single most useful PyTorch skill.

Element-wise Math

import torch

a = torch.tensor([1.0, 2.0, 3.0])
b = torch.tensor([10.0, 20.0, 30.0])
print(a + b, a * b)
print(torch.sqrt(a), a ** 2)

tensor([11., 22., 33.]) tensor([10., 40., 90.])
tensor([1.0000, 1.4142, 1.7321]) tensor([1., 4., 9.])

Matrix Multiplication

Use the @ operator (or torch.matmul). The inner dimensions must match: [2, 3] @ [3, 4] gives [2, 4]. This operation is the heart of every linear layer.

X = torch.ones(2, 3)
W = torch.arange(12, dtype=torch.float32).reshape(3, 4)
print((X @ W).shape)
print(X @ W)

torch.Size([2, 4])
tensor([[12., 15., 18., 21.],
        [12., 15., 18., 21.]])

Note that a * b is element-wise multiplication, not matrix multiplication. Confusing the two is a classic bug.

Reductions Along a Dimension

scores = torch.tensor([[0.1, 0.7, 0.2],
                       [0.5, 0.3, 0.2]])
print(scores.sum())
print(scores.mean(dim=0))
print(scores.argmax(dim=1))
print(scores.max(dim=1, keepdim=True).values.shape)

tensor(2.)
tensor([0.3000, 0.5000, 0.2000])
tensor([1, 0])
torch.Size([2, 1])

dim=0 collapses rows (works down the columns), dim=1 collapses columns (works across each row). argmax(dim=1) on model outputs gives the predicted class for each sample.

Reshaping

x = torch.arange(6)
print(x.reshape(2, 3))
print(x.view(3, -1).shape)
img = torch.zeros(1, 28, 28)
print(img.squeeze(0).shape, img.unsqueeze(0).shape)
print(img.flatten(start_dim=1).shape)

tensor([[0, 1, 2],
        [3, 4, 5]])
torch.Size([3, 2])
torch.Size([28, 28]) torch.Size([1, 1, 28, 28])
torch.Size([1, 784])

-1 means "work out this size for me". view requires contiguous memory and never copies; reshape copies only when it has to, so it is the safer default. permute reorders dimensions, for example from [H, W, C] to [C, H, W].

Combining Tensors

a = torch.zeros(2, 3)
b = torch.ones(2, 3)
print(torch.cat([a, b], dim=0).shape)
print(torch.stack([a, b]).shape)

torch.Size([4, 3])
torch.Size([2, 2, 3])

Broadcasting

When shapes differ, PyTorch stretches size-1 dimensions to match, without copying data. Shapes are compared from the right; each pair must be equal or one of them must be 1.

x = torch.ones(3, 4)
bias = torch.tensor([1.0, 2.0, 3.0, 4.0])   # shape [4]
print(x + bias)
col = torch.tensor([[10.0], [20.0], [30.0]])  # shape [3, 1]
print((x * col).shape)

tensor([[2., 3., 4., 5.],
        [2., 3., 4., 5.],
        [2., 3., 4., 5.]])
torch.Size([3, 4])

Frequently asked questions

What is the difference between view and reshape in PyTorch?

view returns a new view of the same memory and fails on non-contiguous tensors. reshape behaves like view when possible and copies when needed, so it always works.

How do I multiply matrices in PyTorch?

Use A @ B or torch.matmul(A, B). The * operator does element-wise multiplication instead.

What is the difference between torch.cat and torch.stack?

cat joins tensors along an existing dimension. stack creates a new dimension and places the tensors along it, so all inputs must have the same shape.