PyTorch tensors

Every number in PyTorch lives in a tensor. A tensor is a multi-dimensional array with a fixed dtype and a device. Rank 0 is a scalar, rank 1 is a vector and rank 2 is a matrix. Most bugs you will meet come from one of three things: shapes, devices or dtypes.

Shapes, dtypes and devices

Common dtypes:

• float32 is the default for training.
• float16 and bfloat16 are used for mixed precision.
• int64 is used for indices and class labels.
• bool is used for masks.

Devices never move on their own. Move the data explicitly:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x = x.to(device)

Operands must be on the same device. A classic footgun is leaving labels on the CPU while the logits are on the GPU.

Views, copies and strides

A view is a new shape over the same memory. These operations are views or cheap changes:

• view and reshape change the shape.
• transpose and permute reorder the axes.
• contiguous() makes a copy laid out in row order, when you need one.
• clone() always makes a true copy.

a = torch.arange(6).reshape(2, 3)
print(a.stride(), a.t().stride())
print(a.t().is_contiguous())

Output:

(3, 1) (1, 3)
False

A stride says how many steps in memory you take to move one place along an axis. Transposing swaps the strides, so the memory is no longer laid out row by row. That is why view can fail after a transpose. Use reshape, or call .contiguous() first.

Broadcasting

Shapes are compared from the right. Two dimensions match when they are equal or one of them is 1. Here (2, 5, 4) plus (4,) works, because the (4,) vector is repeated across the other axes:

h = torch.randn(2, 5, 4)
bias = torch.tensor([1.0, 2.0, 3.0, 4.0])
print((h + bias).shape)

Output:

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

Broadcasting also hides bugs. Adding a (B,) vector to a (B, T) matrix can silently run along the wrong axis. When something "kind of works," print the shapes.

Common mistakes

• Mixing devices. Move the model, the batch and the targets to the same device.
• Trusting view after a transpose. It needs contiguous memory. Use reshape or .contiguous().
• Broadcasting by accident. Print shape on both sides of a suspicious add or multiply.

Check yourself

1. Why does h.transpose(0, 1).view(10, 4) fail? Answer: the transposed tensor is not contiguous, and view needs contiguous memory. Use reshape instead.
2. What shape does (2, 1, 4) + (3, 1) give? Answer: (2, 3, 4). Align from the right: 4 with 1 gives 4, 1 with 3 gives 3, and 2 stands alone.

Key takeaways

• A tensor is an array with a shape, a dtype and a device.
• Move data to the device explicitly. Mismatched devices raise errors.
• view needs contiguous memory. Strides explain why.
• Broadcasting is convenient and quietly dangerous. Check shapes.

Further learning

Next, learn how PyTorch computes gradients for these tensors in autograd.