A Gentle Introduction to PyTorch: Tensors Before Attention

While working through Sebastian Raschka's Build a Large Language Model (From Scratch), I hit a small but very real obstacle before attention even began: PyTorch tensor syntax. I understood the broad story—tokens become embeddings, embeddings become queries, keys, and values—but I kept stopping on operations such as torch.arange, stack, and cat. The problem was not the mathematics. It was that I had not yet built a reliable picture of what the shapes meant.

So I stopped trying to read Chapter 3 at full speed and started a smaller session: one PyTorch operation, one concrete example, one shape at a time. This is the first set of notes from that reset.

I did not need all of PyTorch

PyTorch is large. The subset I need for the next few chapters is not: construct a tensor, read its shape, combine tensors, add or remove an axis, and understand the difference between element-by-element multiplication and a dot product. That is enough groundwork to make the book's code legible rather than magical.

A tensor is a container of numbers plus a shape. The shape is not incidental metadata; it says how to interpret the axes. A one-dimensional tensor such as tensor([0, 1, 2]) has shape (3,): one axis with three positions.

import torch

token_positions = torch.arange(3)
# tensor([0, 1, 2])
# shape: (3,)

torch.arange(3) starts at zero and stops before three, like Python's range. It is useful when the values follow a regular pattern. When I already know the values, I use torch.tensor(...) instead.

Shapes express relationships

The first useful leap was to treat dimensions as directions. For a two-dimensional tensor with shape (2, 3), dim=0 is the vertical direction: two rows. dim=1 is the horizontal direction: three columns. There is nothing more mystical in the word dimension here.

That makes the distinction between stack and cat memorable:

morning = torch.tensor([10, 20, 30])
evening = torch.tensor([40, 50, 60])

torch.stack([morning, evening]).shape  # (2, 3)
torch.cat([morning, evening]).shape    # (6,)

stack preserves the two inputs as separate rows, creating a new axis. cat places the values end-to-end, extending an axis that already exists. With 2D tensors, the selected dim is the dimension that grows; the remaining dimensions need to match.

Batching is just one extra axis

Language-model examples often look more intimidating than they are because they include a batch axis. A single sequence of three token IDs has shape (3,). Adding an outer axis with unsqueeze(0) produces shape (1, 3): one sequence in the batch, with three token positions.

token_ids = torch.tensor([101, 202, 303])  # (3,)
batch = token_ids.unsqueeze(0)             # (1, 3)

The leading 1 is not an index; it is the batch size. The one sequence happens to be available at batch[0]. Once that distinction clicks, a batch of two sequences is simply shape (2, 3).

Multiplication is not always a dot product

I also needed to separate two operations that are easy to blur together. a * b multiplies matching positions and keeps every result. a @ b multiplies matching positions and sums them, producing one compatibility score for two vectors.

a = torch.tensor([1, 2, 3])
b = torch.tensor([10, 20, 30])

a * b  # tensor([10, 40, 90])
a @ b  # tensor(140)

The second operation is the important precursor to attention: a query and a key will eventually be compared with a dot product. But that is a later lesson. First I want the input and output shapes of ordinary tensor operations to feel obvious.

The stopping point is intentional

This first session is deliberately small. It establishes a concrete vocabulary: values, shapes, axes, rows, columns, batches, concatenation, stacking, and dot products. In the next session I will cover the remaining operations that appear constantly in the book—reshaping, transposing, broadcasting, and softmax—then return to embeddings and attention.

The aim is not to memorise PyTorch calls. It is to be able to look at a line in Raschka's code and answer four questions: what tensor enters, what shape it has, what operation runs, and what shape comes out. Once I can do that, the rest of the transformer has somewhere solid to stand.