Tensors in PyTorch: A Structured Guide for LLM Work

Introduction

PyTorch is the dominant framework for building and training deep learning models. At its core, it provides one fundamental data structure — the tensor — and a rich set of operations over it. Everything else in PyTorch, from neural network layers to optimisers to distributed training, is built on top of tensors.

A tensor is a generalisation of familiar mathematical objects: a scalar is a 0-dimensional tensor, a vector is 1-dimensional, a matrix is 2-dimensional, and anything beyond that is simply called an N-dimensional tensor. What makes PyTorch’s tensors powerful is not just the abstraction, but what comes attached to it: automatic differentiation, GPU acceleration, memory-efficient views, and a broadcasting system that eliminates most explicit loops.

If you are reading model code — whether it is GPT, Llama, or any attention-based architecture — tensors are the language the code is written in. A single forward pass through a transformer is a sequence of roughly 30 tensor operations. Understanding what each one does to the shape, memory, and gradient graph of a tensor is the difference between reading model code and truly understanding it.

This article builds that understanding from the ground up, ending with a fully annotated attention forward pass where every line traces back to a concept explained earlier.

What You Will Learn

By the end of this article you will be able to:

• Create and inspect tensors — understand dtype, device, shape, and numel(), and know when to use factory functions vs torch.tensor()
• Index and slice tensors — use boolean masks, torch.where(), and integer tensor indexing the way transformer code actually uses them
• Manipulate shapes without confusion — know the difference between view() and reshape(), understand why transpose() requires contiguous() before view(), and use squeeze, unsqueeze, expand, and permute correctly
• Understand broadcasting — know the alignment rules, spot silent broadcasting bugs, and see how GQA uses broadcasting to avoid memory copies
• Run math operations efficiently — use reductions with dim= and keepdim=, perform batched matmul with @, read einsum notation, and understand why numerical stability matters for softmax
• Understand memory and strides — know what storage is, why transpose() is zero-copy, what is_contiguous() means, and when clone() vs detach() is the right tool
• Use autograd correctly — understand the computational graph, requires_grad, backward(), no_grad(), and why in-place ops break gradient flow
• Read and write transformer code — follow the head-splitting pattern, understand causal masking with masked_fill(), and trace shapes through a complete attention forward pass
• Debug tensor errors — recognise and fix the five most common runtime errors: device mismatch, matmul shape mismatch, non-contiguous view, dtype mismatch, and silent nan in loss

1. Tensor Basics

What a tensor is

A scalar is a single number. A vector is a list. A matrix is a grid. A tensor is the generalisation — any number of dimensions.

scalar = torch.tensor(3.14)          # shape: []        — 0 dimensions
vector = torch.tensor([1, 2, 3])     # shape: [3]       — 1 dimension
matrix = torch.tensor([[1,2],[3,4]]) # shape: [2, 2]    — 2 dimensions
cube   = torch.zeros(2, 3, 4)        # shape: [2, 3, 4] — 3 dimensions

In LLM work you constantly deal with 4D tensors: [batch, heads, seq_len, head_dim]. That's just a tensor with 4 dimensions — the math is identical to the 2D case.

torch.tensor() vs torch.Tensor() vs factory functions

# torch.tensor() — copies data, infers dtype from input
a = torch.tensor([1.0, 2.0])     # float32 (Python float → float32)
b = torch.tensor([1, 2])         # int64   (Python int → int64)

# torch.Tensor() — always float32, no dtype inference
c = torch.Tensor([1, 2])         # float32 regardless

# Factory functions — preferred in practice
torch.zeros(2, 3)
# tensor([[0., 0., 0.],
#         [0., 0., 0.]])

torch.ones(2, 3)
# tensor([[1., 1., 1.],
#         [1., 1., 1.]])

torch.arange(0, 10, 2)
# tensor([0, 2, 4, 6, 8])

torch.linspace(0, 1, 5)
# tensor([0.0000, 0.2500, 0.5000, 0.7500, 1.0000])

Rule of thumb: use torch.tensor() when you have existing data, factory functions when you need a fresh tensor.

torch.empty() — allocate without initialising

torch.empty(*shape) allocates a tensor of the given shape without writing any values into it. The contents are whatever bytes happened to already be in that memory — garbage, not zeros.

torch.empty(2, 3)
# tensor([[ 1.4013e-45,  0.0000e+00, -3.4194e+13],
#         [ 4.5648e-41,  8.9683e-44,  0.0000e+00]])   — arbitrary leftover memory, different every run

It’s faster than torch.zeros() or torch.ones() because it skips the initialisation pass — worth it only when you are about to overwrite every element yourself and don't want to pay for a write you'll immediately discard.

# Pre-allocate a KV cache buffer, then fill it incrementally with scatter_
kv_cache = torch.empty(1, 8, 2048, 64)   # allocate once
kv_cache[:, :, :5, :] = new_kv           # write real values before ever reading the rest

Danger: reading from an empty tensor before writing to every position it will be used from is a silent correctness bug — you get real numbers, not an error, and they will be different each run. Never use torch.empty() when you need a zero baseline (use torch.zeros() or torch.full_like() instead); reserve it for buffers that get fully overwritten before use.

torch.full_like() and the _like factory functions

torch.full_like(x, value) creates a tensor filled with value, matching x's shape, dtype, and device. It's the fastest way to build a constant tensor that's guaranteed to line up with an existing one.

x = torch.randn(2, 3, dtype=torch.bfloat16, device="cuda")

torch.full_like(x, -1e9)
# tensor([[-1e9, -1e9, -1e9],
#         [-1e9, -1e9, -1e9]], dtype=torch.bfloat16, device='cuda:0')
# same shape, dtype, and device as x — no manual bookkeeping

torch.zeros_like(x)   # same as torch.full_like(x, 0)
torch.ones_like(x)    # same as torch.full_like(x, 1)

This shows up when building an additive attention bias or a padding fill value that has to exactly match the score tensor’s shape and dtype:

scores = torch.randn(2, 8, 512, 512)
neg_inf_bias = torch.full_like(scores, float('-inf'))
scores = torch.where(causal_mask, neg_inf_bias, scores)

Using torch.full_like() instead of torch.full(scores.shape, ...) avoids a whole class of dtype/device-mismatch bugs (section 12) — you never have to remember to also pass dtype= and device=.

token_ids = torch.tensor([101, 2003, 1037], dtype=torch.int64)
mask      = torch.tensor([True, True, False], dtype=torch.bool)
weights   = torch.randn(768, 768, dtype=torch.bfloat16)

bfloat16 has the same exponent range as float32 (so it doesn't overflow) but half the precision — that's why it's preferred for training over float16.

device

A tensor lives on exactly one device. Operations between tensors on different devices fail.

cpu_tensor  = torch.randn(3, 3)                        # on CPU
gpu_tensor  = torch.randn(3, 3, device="cuda")         # on GPU
gpu_tensor2 = torch.randn(3, 3, device="cuda:1")       # on second GPU

# This will error:
cpu_tensor + gpu_tensor   # RuntimeError: expected all tensors on same device

shape, ndim, numel()

x = torch.randn(2, 8, 512, 64)  # batch=2, heads=8, seq=512, head_dim=64

x.shape      # torch.Size([2, 8, 512, 64])
x.ndim       # 4
x.numel()    # 2 × 8 × 512 × 64 = 524288 — total number of elements

2. Indexing and Slicing

Basic slicing

x = torch.tensor([[1, 2, 3],
                  [4, 5, 6],
                  [7, 8, 9]])

x[0]        # tensor([1, 2, 3])         — first row
x[1:3]      # tensor([[4,5,6],[7,8,9]]) — rows 1 and 2
x[:, 0]     # tensor([1, 4, 7])         — first column, all rows
x[1, 2]     # tensor(6)                 — row 1, col 2

The : means "all of this dimension". [:, 0] reads as "every row, column 0".

Boolean masking

x = torch.tensor([3, -1, 4, -1, 5])

mask = x > 0
# tensor([ True, False,  True, False,  True])

x[mask]
# tensor([3, 4, 5])    — only the elements where mask is True

x[x < 0] = 0          # in-place: zero out negatives
# tensor([3, 0, 4, 0, 5])

torch.where()

torch.where(condition, x, y) picks from x where condition is True, from y where False.

scores = torch.tensor([0.8, -1e9, 0.3, -1e9, 0.5])
mask   = torch.tensor([True, False, True, False, True])

# Replace masked positions with -inf before softmax (causal masking)
result = torch.where(mask, scores, torch.tensor(float('-inf')))
# tensor([0.8, -inf, 0.3, -inf, 0.5])

This is exactly how causal attention masks work — positions the token cannot attend to become -inf, so softmax drives them to zero.

Advanced indexing with integer tensors

vocab = torch.randn(50000, 768)   # embedding table: 50k tokens, dim 768
token_ids = torch.tensor([101, 2003, 1037])

# Index with a tensor of indices — retrieves 3 rows
embeddings = vocab[token_ids]     # shape: [3, 768]

This is the embedding lookup every transformer does on input token IDs.

torch.scatter_() — writing values at indexed positions

Where x[token_ids] reads rows by index, scatter_ writes values into a tensor at positions given by an index tensor. x.scatter_(dim, index, src) writes, for every position, x[..., index[...], ...] = src[...] along dim.

x = torch.zeros(3, 5)
index = torch.tensor([[1], [0], [4]])
src = torch.tensor([[1.0], [1.0], [1.0]])

x.scatter_(1, index, src)
# tensor([[0., 1., 0., 0., 0.],
#         [1., 0., 0., 0., 0.],
#         [0., 0., 0., 0., 1.]])
# row 0 got a 1 at column 1, row 1 at column 0, row 2 at column 4

The classic use is building one-hot labels from class indices:

labels = torch.tensor([2, 0, 1])          # class index per example
one_hot = torch.zeros(3, 3)
one_hot.scatter_(1, labels.unsqueeze(1), 1.0)
# tensor([[0., 0., 1.],
#         [1., 0., 0.],
#         [0., 1., 0.]])

It also shows up in KV-cache updates, where a new token’s key/value vectors are written into a fixed-size cache at the current position:

kv_cache = torch.zeros(1, 8, 2048, 64)          # [batch, heads, max_seq, head_dim]
position  = torch.tensor([5]).view(1, 1, 1, 1).expand(1, 8, 1, 64)
new_kv    = torch.randn(1, 8, 1, 64)

kv_cache.scatter_(2, position, new_kv)   # writes new_kv at seq position 5

The trailing underscore means this is an in-place op — same caution applies as with any _ method (section 5): it mutates x directly and will error on a leaf tensor that requires_grad. torch.scatter() (no underscore) returns a new tensor instead of mutating.

View (shared memory) vs copy

x = torch.arange(6)
# tensor([0, 1, 2, 3, 4, 5])

y = x.view(2, 3)
# tensor([[0, 1, 2],
#         [3, 4, 5]])          — same data, different shape

y[0, 0] = 99
print(x)
# tensor([99,  1,  2,  3,  4,  5])  — x changed because y shares x's memory

z = x.clone().view(2, 3)     # clone first → independent copy
z[0, 0] = 0
print(x)                     # x is unchanged

Most shape operations (view, transpose, slices) return views. Surprising mutations are the classic bug when you forget this.

3. Shape Manipulation

This section is the most important for LLM work. Almost every transformer operation is a sequence of reshapes.

view() vs reshape()

Both change shape without moving data — when possible. The difference is what happens when it’s not possible.

x = torch.arange(12)
# tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11])

x.view(3, 4)
# tensor([[ 0,  1,  2,  3],
#         [ 4,  5,  6,  7],
#         [ 8,  9, 10, 11]])    — rows are filled left-to-right from the flat list

x.view(4, 3)
# tensor([[ 0,  1,  2],
#         [ 3,  4,  5],
#         [ 6,  7,  8],
#         [ 9, 10, 11]])

x.reshape(3, 4)   # same output as view — but if a view isn't possible, copies

After a transpose(), the tensor is no longer contiguous in memory (more on this in section 6). view() will error; reshape() will silently copy.

x = torch.randn(3, 4)
t = x.transpose(0, 1)    # shape [4, 3], but memory layout is still [3, 4]

t.view(12)               # RuntimeError: not contiguous
t.reshape(12)            # works — makes a copy internally
t.contiguous().view(12)  # explicit: make contiguous, then view

Rule: use view() intentionally when you know the tensor is contiguous and you want a guaranteed zero-copy operation. Use reshape() when you don't need that guarantee.

transpose() and permute()

transpose() swaps exactly two dimensions. permute() reorders all dimensions at once.

# Small example to see what the data looks like after transposing
m = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])    # shape: [2, 3]

m.transpose(0, 1)
# tensor([[1, 4],
#         [2, 5],
#         [3, 6]])               shape: [3, 2] — rows became columns

For higher-dimensional tensors (the LLM case), only the shapes change — the pattern is the same:

x = torch.randn(2, 8, 512, 64)  # [batch, heads, seq, head_dim]

x.transpose(2, 3)      # shape: [2, 8, 64, 512]  — seq and head_dim swapped
x.permute(0, 2, 1, 3)  # shape: [2, 512, 8, 64]  — heads moved to dim 2

squeeze() and unsqueeze()

unsqueeze(dim) adds a dimension of size 1. squeeze(dim) removes it.

x = torch.tensor([1, 2, 3])   # shape: [3]
x.unsqueeze(0)
# tensor([[1, 2, 3]])          shape: [1, 3] - batch dim added in front
x.unsqueeze(1)
# tensor([[1],
#         [2],
#         [3]])                shape: [3, 1] - new dim inserted after dim 0
y = torch.tensor([[[1, 2, 3]]])   # shape: [1, 1, 3]
y.squeeze()
# tensor([1, 2, 3])           shape: [3] - all size-1 dims removed
y.squeeze(0)
# tensor([[1, 2, 3]])         shape: [1, 3] - only dim 0 removed

unsqueeze is everywhere in broadcasting setups — you add a dim so PyTorch can broadcast across it.

expand() vs repeat()

x = torch.tensor([[1], [2], [3]])   # shape: [3, 1]
# tensor([[1],
#         [2],
#         [3]])
x.expand(3, 4)
# tensor([[1, 1, 1, 1],
#         [2, 2, 2, 2],
#         [3, 3, 3, 3]])    shape: [3, 4] - zero copy, just changes strides
x.repeat(1, 4)
# tensor([[1, 1, 1, 1],
#         [2, 2, 2, 2],
#         [3, 3, 3, 3]])    shape: [3, 4] - allocates new memory, copies data

Both produce the same visible result — the difference is only in memory: expand() reads the same element repeatedly via stride tricks, repeat() physically writes four copies.

expand() is preferred when you're about to use the result in a computation. repeat() is for when you genuinely need a concrete copy — rare.

In GQA (Grouped Query Attention), K and V have fewer heads than Q. expand() broadcasts K and V to match Q's head count without allocating memory for the repeated heads:

# Q: [batch, 32, seq, head_dim]  (32 query heads)
# K: [batch,  4, seq, head_dim]  (4 KV heads, grouped)
k = k.unsqueeze(2)                       # [batch, 4, 1, seq, head_dim]
k = k.expand(-1, -1, 8, -1, -1)         # [batch, 4, 8, seq, head_dim]
k = k.reshape(batch, 32, seq, head_dim) # merge group dim back

contiguous()

After transpose() or permute(), the tensor's memory layout doesn't match its shape. contiguous() makes a copy that does.

x = torch.randn(3, 4)
t = x.transpose(0, 1)
t.is_contiguous()               # False
t.contiguous().is_contiguous()  # True

You need this before view() — see section 6 for why.

flatten() and unflatten()

x = torch.tensor([[[1, 2], [3, 4]],
                  [[5, 6], [7, 8]]])   # shape: [2, 2, 2]
x.flatten()
# tensor([1, 2, 3, 4, 5, 6, 7, 8])    shape: [8] - everything into one list
x.flatten(1)
# tensor([[1, 2, 3, 4],
#         [5, 6, 7, 8]])               shape: [2, 4] - keep dim 0, flatten rest
x.flatten(0, 1)
# tensor([[1, 2],
#         [3, 4],
#         [5, 6],
#         [7, 8]])                     shape: [4, 2] - flatten only dims 0 and 1

For large LLM tensors the shapes work the same way — only the numbers differ:

x = torch.randn(2, 8, 512, 64)
x.flatten(2, 3)           # [2, 8, 32768] — merge seq and head_dim
y = torch.randn(2, 32768)
y.unflatten(1, (512, 64)) # [2, 512, 64]  - split dim 1 back into (seq, head_dim)

4. Broadcasting

The rules

PyTorch aligns shapes from the right, then expands any dimension that is 1.

Shape A:    [   8, 512,  64]
Shape B:    [      1,   64]   ← aligned from right
Result:     [   8, 512,  64]  ← B's size-1 dims expand to match A

A concrete example with small numbers to see what actually happens:

a = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])    # shape: [2, 3]
b = torch.tensor([[10],
                  [20]])         # shape: [2, 1]
a + b
# tensor([[11, 12, 13],
#         [24, 25, 26]])
# b's column [10, 20] was virtually repeated 3 times to match a's 3 columns

In LLM code the same rule applies to larger shapes:

scores = torch.randn(8, 512, 512)   # [heads, seq, seq]
mask   = torch.zeros(1, 512, 512)   # [1, seq, seq]
scores + mask   # mask broadcasts across the 8 heads - no data copy

When broadcasting copies data vs when it’s zero-copy

expand() is zero-copy — it just adjusts strides so the same memory element is read multiple times. But once you pass a broadcasted tensor into most ops (like + or matmul), PyTorch will materialise the expanded tensor to perform the computation. The view is zero-copy; the result of the op is not.

Common broadcasting bugs

a = torch.randn(512, 64)
b = torch.randn(64, 512)
a + b   # RuntimeError - shapes don't align from the right
        # Right-align: [512, 64] vs [64, 512] → last dims 64 ≠ 512

The silent version is worse — shapes that almost match:

a = torch.randn(8, 1, 512)
b = torch.randn(1, 512, 8)
a + b   # shape: [8, 512, 8] - no error, but probably not what you wanted

Always check .shape after operations when debugging unexpected results.

How GQA uses broadcasting

In Grouped Query Attention, K and V heads are shared across groups of Q heads. Broadcasting expands them without copying:

# Q: [batch, 32, seq, head_dim]
# K: [batch,  4, seq, head_dim]
k = k.unsqueeze(2).expand(-1, -1, 8, -1, -1)
# K is now virtually [batch, 4, 8, seq, head_dim] - same memory, different strides

5. Math Operations

Elementwise

x = torch.tensor([1.0, 4.0, 9.0])
x + 1          # [2., 5., 10.]
x ** 2         # [1., 16., 81.]
torch.sqrt(x)  # [1., 2., 3.]
torch.exp(x)   # [e¹, e⁴, e⁹]

Reductions — dim= and keepdim=

x = torch.tensor([[1., 2., 3.],
                  [4., 5., 6.]])   # shape: [2, 3]
x.sum()
# tensor(21.)                              - all 6 elements summed
x.sum(dim=0)
# tensor([5., 7., 9.])                     shape: [3] - each column summed across rows
x.sum(dim=1)
# tensor([ 6., 15.])                       shape: [2] - each row summed across columns
x.sum(dim=1, keepdim=True)
# tensor([[ 6.],
#         [15.]])                           shape: [2, 1] - dim 1 kept as size-1

keepdim=True matters for broadcasting — without it, the reduced tensor loses a dimension and may not align correctly in subsequent operations.

# Softmax by hand (why keepdim matters)
x = torch.randn(4, 512)
x_max = x.max(dim=1, keepdim=True).values   # [4, 1] — broadcasts correctly
x = x - x_max                               # [4, 512] - [4, 1] → fine

torch.topk

torch.topk(x, k, dim=) returns the k largest values along dim, plus their indices. It returns a named tuple (values, indices).

x = torch.tensor([3.0, 1.0, 4.0, 1.0, 5.0, 9.0, 2.0])
torch.topk(x, 3)
# torch.return_types.topk(
# values=tensor([9., 5., 4.]),
# indices=tensor([5, 4, 2]))
# largest 3 values, in descending order, with their original positions

The signature LLM use is top-k sampling during generation — restrict the next-token distribution to the k highest-probability tokens before sampling:

logits = torch.randn(1, 50000)     # [batch, vocab]
k = 50
top_values, top_indices = torch.topk(logits, k, dim=-1)   # both [1, 50]
# Mask everything outside the top-k before softmax
filtered = torch.full_like(logits, float('-inf'))
filtered.scatter_(-1, top_indices, top_values)
probs = torch.softmax(filtered, dim=-1)   # [1, 50000], zero everywhere but the top-k
next_token = torch.multinomial(probs, num_samples=1)

This is also how a router in a Mixture-of-Experts layer picks which experts to activate per token: torch.topk(gate_logits, num_experts_per_token, dim=-1).

largest=False flips it to smallest-k; sorted=False skips the sort for a small speedup when order doesn't matter.

Matrix operations

a = torch.randn(512, 64)
b = torch.randn(64, 128)
torch.matmul(a, b)   # [512, 128]
a @ b                # same - @ is the matmul operator
# Batched matmul - @ works on 3D/4D tensors too
q = torch.randn(2, 8, 512, 64)   # [batch, heads, seq, head_dim]
k = torch.randn(2, 8, 64, 512)   # [batch, heads, head_dim, seq]
scores = q @ k    # [2, 8, 512, 512] - matmul over last two dims, batch over first two

torch.bmm() is the older batched matmul — only handles 3D tensors. The @ operator is more general and preferred.

torch.bmm() — strict batched matmul

torch.bmm(a, b) multiplies a batch of matrices: a is [B, M, K], b is [B, K, N], result is [B, M, N]. Unlike @, it requires exactly 3D inputs and does no broadcasting over the batch dimension — the batch sizes must match exactly.

a = torch.randn(4, 512, 64)
b = torch.randn(4, 64, 128)
torch.bmm(a, b)   # [4, 512, 128]
# No broadcasting allowed:
c = torch.randn(1, 64, 128)
torch.bmm(a, c)   # RuntimeError: batch1 and batch2 must have same batch size
a @ c             # works - [4,512,64] @ [1,64,128] broadcasts to [4,512,128]

For 4D attention tensors ([batch, heads, seq, head_dim]) you'd first collapse batch and heads into one dim to use bmm:

q = torch.randn(2, 8, 512, 64)
k = torch.randn(2, 8, 64, 512)
q_flat = q.flatten(0, 1)                 # [16, 512, 64]
k_flat = k.flatten(0, 1)                 # [16, 64, 512]
scores = torch.bmm(q_flat, k_flat)       # [16, 512, 512]
scores = scores.unflatten(0, (2, 8))     # [2, 8, 512, 512]

In practice @ does this for you without the manual flatten/unflatten, which is why bmm mostly appears in older codebases or when you want the stricter shape check as a safety net against silent broadcasting bugs (section 4).

einsum

Einstein summation — expresses any contraction or reordering in one string.

# Attention: scores[b,h,s,S] = sum over d of q[b,h,s,d] * k[b,h,S,d]
scores = torch.einsum("bhsd,bhSd->bhsS", q, k)
# Outer product
a = torch.tensor([1., 2., 3.])
b = torch.tensor([4., 5., 6., 7.])
torch.einsum("i,j->ij", a, b)
# tensor([[ 4.,  5.,  6.,  7.],
#         [ 8., 10., 12., 14.],
#         [12., 15., 18., 21.]])   shape: [3, 4] - every pair multiplied
# Batch matrix multiply
torch.einsum("bik,bkj->bij", x, y)   # equivalent to x @ y for 3D

Read the string as: name the dims of each input, name the dims of the output. Any dim that appears in inputs but not the output gets summed over.

Softmax and numerical stability

x = torch.tensor([1.0, 2.0, 3.0])
torch.softmax(x, dim=0)    # [0.09, 0.24, 0.67]

Why numerical stability matters:

x = torch.tensor([1000.0, 1001.0, 1002.0])
torch.softmax(x, dim=0)            # works — PyTorch subtracts max internally
# Naive implementation:
torch.exp(x) / torch.exp(x).sum()  # inf/inf → nan

PyTorch’s softmax subtracts max(x) before exponentiating — this doesn't change the result mathematically but prevents overflow. log_softmax is additionally preferred for loss computation because log(softmax(x)) has worse numerical properties than computing log_softmax directly.

In-place ops and autograd

x = torch.tensor([1.0, 2.0], requires_grad=True)
x.add_(1)   # in-place - modifies x directly
# RuntimeError: a leaf Variable that requires grad has been used in an in-place operation

PyTorch’s autograd records operations to compute gradients. In-place ops destroy the original value that the backward pass needs to use. Avoid them on any tensor that’s part of a computation graph.

6. Memory and Storage

Storage and strides

Every tensor is a view into a flat 1D block of memory called storage. The tensor’s strides tell PyTorch how many elements to skip in storage to advance one step in each dimension.

x = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])   # shape: [2, 3]
list(x.storage())   # [1, 2, 3, 4, 5, 6] - one flat block in memory
x.stride()          # (3, 1)
#                     ↑  ↑
#                     |  └─ move 1 element to step right one column
#                     └──── move 3 elements to step down one row

To find element x[i, j]: storage[i * 3 + j * 1].

So x[1, 2] = storage[1*3 + 2*1] = storage[5] = 6. ✓

Why transpose() doesn't move data

t = x.transpose(0, 1)   # shape: [3, 2]
list(t.storage())   # [1, 2, 3, 4, 5, 6] - identical to x, nothing moved
t.stride()          # (1, 3) - strides are just swapped
print(t)
# tensor([[1, 4],
#         [2, 5],
#         [3, 6]])

t[i, j] = storage[i * 1 + j * 3]. So t[0, 1] = storage[0 + 3] = storage[3] = 4. The same flat memory, read with different navigation rules.

This is why transpose() is cheap and why view() fails afterward: view() requires elements to be consecutive in storage so it can slice them into rows, and after transposing they are interleaved instead.

is_contiguous()

x = torch.randn(3, 4)
x.is_contiguous()               # True — row-major, as expected
t = x.transpose(0, 1)
t.is_contiguous()               # False - strides no longer match shape
t.contiguous().is_contiguous()  # True - made a fresh copy in memory

clone() vs detach() vs detach().clone()

x = torch.randn(3, requires_grad=True)
y = x * 2   # y is in the computation graph
y.clone()            # new tensor, same graph - gradient still flows through y
y.detach()           # same storage as y, removed from graph - no copy
y.detach().clone()   # new tensor, removed from graph - fully independent

Use .detach().clone() when you want a plain tensor you can inspect or log, with no graph attachment and no shared memory.

7. Autograd — How Gradients Flow Through Tensors

requires_grad=True

x = torch.tensor([2.0], requires_grad=True)
y = x ** 2 + 3 * x + 1   # y = x² + 3x + 1
y.backward()
x.grad   # tensor([7.]) - dy/dx = 2x + 3 = 2(2) + 3 = 7

Only leaf tensors (those you created, not computed) accumulate .grad. Intermediate tensors don't by default.

The computational graph

PyTorch builds a graph dynamically as you do operations. Each tensor stores a reference to the function that created it.

x = torch.randn(3, requires_grad=True)
y = x.sum()
y.grad_fn                  # <SumBackward0>
y.grad_fn.next_functions   # points back to x

backward() walks this graph in reverse, applying the chain rule at each node.

torch.no_grad()

with torch.no_grad():
    output = model(input)   # no graph built — saves memory and compute

During inference you don’t need gradients. no_grad() prevents PyTorch from building the graph, which reduces memory usage and speeds up the forward pass.

detach()

# Stop gradient flowing through a particular path
target = output.detach()   # treat as a constant, not a trainable output
loss = F.mse_loss(prediction, target)

Common in RL and contrastive learning where you want one branch of a computation to not receive gradients.

retain_graph=True

loss.backward(retain_graph=True)   # keep the graph alive
loss.backward()                    # use it again

By default, PyTorch frees the graph after backward() to save memory. retain_graph=True keeps it — needed when you call backward multiple times (e.g. multiple loss terms or meta-learning).

8. Type and Device Movement

x = torch.randn(3, 3)
x.to("cuda")                          # move to GPU
x.to("cpu")                           # move to CPU
x.to(torch.bfloat16)                  # cast dtype
x.to("cuda", dtype=torch.bfloat16)   # both at once - most efficient, one copy
# Shorthands
x.cuda()      # to GPU
x.cpu()       # to CPU
x.float()     # float32
x.half()      # float16
x.bfloat16()  # bfloat16

Common bug: two tensors on different devices or with different dtypes error at the operation, not at creation. The fix is to cast/move before the op:

a = torch.randn(3, device="cpu")
b = torch.randn(3, device="cuda")
a + b               # RuntimeError: expected all on same device
# Fix:
a.to(b.device) + b

9. Tensor Operations Critical for Transformer Internals

Everything in sections 1–8 feeds into these patterns.

Batched matmul for attention scores

q = torch.randn(2, 8, 512, 64)   # [batch, heads, seq_q, head_dim]
k = torch.randn(2, 8, 512, 64)   # [batch, heads, seq_k, head_dim]
# Compute all attention scores in one op
scores = q @ k.transpose(-2, -1)  # [2, 8, 512, 512]
# transpose(-2, -1) swaps the last two dims: [batch, heads, head_dim, seq_k]

@ on 4D tensors batches over all leading dims and does matmul over the last two.

view() + transpose() head-splitting pattern

This is the pattern in every transformer implementation:

batch, seq_len = 2, 512
num_heads, head_dim = 8, 64
hidden_dim = num_heads * head_dim   # 512
# After linear projection: flat hidden dim
q = torch.randn(batch, seq_len, hidden_dim)  # [2, 512, 512]
# Split into heads
q = q.view(batch, seq_len, num_heads, head_dim)  # [2, 512, 8, 64]
q = q.transpose(1, 2)                            # [2, 8, 512, 64]

view() reinterprets the last 512 elements as 8 groups of 64 (no data movement). transpose() reorders dims so each head's queries are grouped together.

torch.cat() vs torch.stack()

a = torch.randn(4, 64)
b = torch.randn(4, 64)
torch.cat([a, b], dim=0)    # [8, 64]   - join along existing dim 0
torch.cat([a, b], dim=1)    # [4, 128]  - join along existing dim 1
torch.stack([a, b], dim=0)  # [2, 4, 64] - new dim 0
torch.stack([a, b], dim=1)  # [4, 2, 64] - new dim 1

cat joins along an existing dimension; stack creates a new one. Use stack when assembling a batch from individual items, cat when appending to an existing dimension.

torch.split() and torch.chunk()

x = torch.randn(4, 12)
# split into pieces of specified size
torch.split(x, 4, dim=1)    # three tensors of [4, 4]
# split into N equal chunks
torch.chunk(x, 3, dim=1)    # three tensors of [4, 4]

In QKV projections, a single linear layer outputs all three:

qkv = linear(x)                              # [batch, seq, 3 * hidden]
q, k, v = torch.split(qkv, hidden, dim=-1)  # each [batch, seq, hidden]

einsum for attention

q = torch.randn(2, 8, 512, 64)
k = torch.randn(2, 8, 512, 64)
v = torch.randn(2, 8, 512, 64)
# Attention scores - d is summed over
scores = torch.einsum("bhsd,bhSd->bhsS", q, k)   # [2, 8, 512, 512]
# b=batch, h=head, s=query_seq, S=key_seq, d=head_dim
attn = torch.softmax(scores / 64**0.5, dim=-1)
# Weighted sum of values
out = torch.einsum("bhsS,bhSd->bhsd", attn, v)   # [2, 8, 512, 64]

Read the string as: name the dims of each input, name the dims of the output. Any dim that appears in inputs but not the output gets summed over.

masked_fill() for causal masks

# With seq_len=4 to make the mask visible
mask = torch.triu(torch.ones(4, 4), diagonal=1).bool()
# tensor([[False,  True,  True,  True],
#         [False, False,  True,  True],
#         [False, False, False,  True],
#         [False, False, False, False]])
#
# Row i = token i. True = "this position is masked" (future token).
# Token 0 can only attend to itself.
# Token 3 can attend to tokens 0, 1, 2, 3.
scores = torch.zeros(1, 1, 4, 4)   # simplified scores
scores = scores.masked_fill(mask, float('-inf'))
# tensor([[[[0., -inf, -inf, -inf],
#           [0.,  0.,  -inf, -inf],
#           [0.,  0.,   0.,  -inf],
#           [0.,  0.,   0.,   0.]]]])
torch.softmax(scores, dim=-1)
# tensor([[[[1.0000, 0.0000, 0.0000, 0.0000],
#           [0.5000, 0.5000, 0.0000, 0.0000],
#           [0.3333, 0.3333, 0.3333, 0.0000],
#           [0.2500, 0.2500, 0.2500, 0.2500]]]])
#
# -inf → 0.0 after softmax. Each row sums to 1.
# Token 0 attends only to itself; token 3 attends equally to all four tokens.

For the full-sequence case, only the shape differs — the masking pattern is the same:

scores = torch.randn(4, 512, 512)   # [batch, seq, seq]
mask   = torch.triu(torch.ones(512, 512), diagonal=1).bool()
scores = scores.masked_fill(mask, float('-inf'))
attn   = torch.softmax(scores, dim=-1)

10. Performance and Memory Efficiency

Tracking VRAM

print(torch.cuda.memory_allocated() / 1e9, "GB")     # currently in use
print(torch.cuda.max_memory_allocated() / 1e9, "GB") # peak since last reset
torch.cuda.reset_peak_memory_stats()

Gradient checkpointing

Normally PyTorch saves all intermediate activations for the backward pass. At large sequence lengths or batch sizes, this exhausts GPU memory.

Gradient checkpointing discards activations during the forward pass and recomputes them during backward — trading ~30% extra compute for a large memory saving.

from torch.utils.checkpoint import checkpoint
output = checkpoint(transformer_block, hidden_states)

pin_memory=True

loader = DataLoader(dataset, batch_size=32, pin_memory=True)

Pinned memory is page-locked on CPU — the GPU can fetch it directly via DMA without an extra CPU copy. Makes CPU→GPU transfer noticeably faster when the GPU is the bottleneck.

Tensor parallelism reshaping patterns

In column-parallel linear (split weight across GPUs by output columns):

# Each GPU holds a slice of the output dim
weight_slice = full_weight[:, start:end]   # [in, out/N]
output_slice = input @ weight_slice        # [batch, seq, out/N]
# All-gather across GPUs to reconstruct full output

In row-parallel linear (split by input dim):

input_slice  = full_input[:, :, start:end]  # [batch, seq, in/N]
output_slice = input_slice @ weight_slice   # partial sum
# All-reduce across GPUs to sum partial results

The reshaping is just a view() or slice on the weight tensor — the parallelism comes from distributing those slices across devices.

torch.compile()

model = torch.compile(model)   # JIT-traces the model, fuses ops

torch.compile() (introduced in PyTorch 2.0) traces your model and applies kernel fusion — for example, merging the matmul + scale + softmax in attention into a single kernel. The biggest gains are in models with many small ops that are individually memory-bandwidth-bound.

The Learning Path

The sections connect like this:

• Basics (1) + Indexing (2) give you the vocabulary.
• Shape manipulation (3) is where the Llama patterns live — view() + transpose() makes sense once you see that they're just adjusting strides over a flat storage block.
• Memory and storage (6) explains why contiguity matters after transpose and why view is zero-copy.
• Broadcasting (4) explains how masks and GQA head expansion work without extra memory.
• Autograd (7) explains why in-place ops and detach patterns exist.
• Transformer internals (9) is all of the above applied — every line in modeling_llama.py maps back to a concept from sections 1–8.

11. The Full Attention Forward Pass — Annotated

Every concept from sections 1–10 appears somewhere in these ~30 lines. Read the shape comments as a running trace of what the tensor looks like at each step.

import torch
import torch.nn.functional as F
# ── Hyperparameters ──────────────────────────────────────────────────────────
batch      = 2
seq_len    = 512
hidden_dim = 512
num_heads  = 8
head_dim   = hidden_dim // num_heads   # 64
# ── Inputs ───────────────────────────────────────────────────────────────────
# Token IDs from the tokeniser
token_ids = torch.randint(0, 50000, (batch, seq_len))          # [2, 512]  int64
# Embedding lookup - each token ID maps to a learned vector
embedding_table = torch.randn(50000, hidden_dim)               # [50000, 512]
x = embedding_table[token_ids]                                 # [2, 512, 512]
# ── Linear projections (Q, K, V) ─────────────────────────────────────────────
# In practice these are nn.Linear layers; here we use raw weight matrices
W_q = torch.randn(hidden_dim, hidden_dim)                      # [512, 512]
W_k = torch.randn(hidden_dim, hidden_dim)
W_v = torch.randn(hidden_dim, hidden_dim)
q = x @ W_q    # [2, 512, 512] @ [512, 512] → [2, 512, 512]
k = x @ W_k    # [2, 512, 512]
v = x @ W_v    # [2, 512, 512]
# ── Split into heads ──────────────────────────────────────────────────────────
# view() splits the last dim (512) into (num_heads=8, head_dim=64) - zero copy
# transpose(1, 2) moves heads forward - changes strides, no data movement
q = q.view(batch, seq_len, num_heads, head_dim).transpose(1, 2)  # [2, 8, 512, 64]
k = k.view(batch, seq_len, num_heads, head_dim).transpose(1, 2)  # [2, 8, 512, 64]
v = v.view(batch, seq_len, num_heads, head_dim).transpose(1, 2)  # [2, 8, 512, 64]
# ── Attention scores ──────────────────────────────────────────────────────────
# k.transpose(-2, -1) → [2, 8, 64, 512]
# @ batches over [batch, heads], matmuls over [seq, head_dim] × [head_dim, seq]
scale  = head_dim ** 0.5                                          # 8.0
scores = (q @ k.transpose(-2, -1)) / scale                       # [2, 8, 512, 512]
# ── Causal mask ───────────────────────────────────────────────────────────────
# Upper triangle is True - position i must not attend to j > i
causal_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
                                                                  # [512, 512]
# masked_fill broadcasts causal_mask across [batch, heads] dims
scores = scores.masked_fill(causal_mask, float('-inf'))           # [2, 8, 512, 512]
# ── Softmax ───────────────────────────────────────────────────────────────────
# dim=-1 normalises over the key dimension (last dim)
# -inf positions become 0 after softmax - those tokens are ignored
attn_weights = torch.softmax(scores, dim=-1)                      # [2, 8, 512, 512]
# ── Weighted sum of values ────────────────────────────────────────────────────
# attn_weights: [2, 8, 512, 512] @ v: [2, 8, 512, 64] → [2, 8, 512, 64]
attn_output = attn_weights @ v                                    # [2, 8, 512, 64]
# ── Merge heads back ──────────────────────────────────────────────────────────
# contiguous() required because transpose() changed strides - view() needs contiguous layout
# view() merges (num_heads, head_dim) back into hidden_dim
attn_output = attn_output.transpose(1, 2).contiguous()           # [2, 512, 8, 64]
attn_output = attn_output.view(batch, seq_len, hidden_dim)       # [2, 512, 512]
# ── Output projection ─────────────────────────────────────────────────────────
W_o = torch.randn(hidden_dim, hidden_dim)
output = attn_output @ W_o                                        # [2, 512, 512]

Summary

Tensors are the single data structure underlying all of PyTorch. Everything else — layers, optimisers, autograd, distributed training — operates on tensors and returns tensors.

The concepts in this article form a dependency chain:

• Storage and strides (section 6) explain why view() is zero-copy and why transpose() breaks it. Without this, the behaviour of shape operations looks arbitrary.
• Shape manipulation (section 3) is the vocabulary of transformer code. The view() + transpose() head-splitting pattern, the contiguous() requirement, expand() vs repeat() — these appear in every attention implementation.
• Broadcasting (section 4) is how masks and GQA head expansion work efficiently. Knowing the right-alignment rule lets you read shape errors immediately instead of guessing.
• Autograd (section 7) is what makes PyTorch a training framework rather than a numerical library. Understanding the computational graph, requires_grad, and detach() is required for writing anything that trains correctly.
• The full attention pass (section 11) is the payoff. Every line in that example traces back to one of the concepts above. Once you can read that code with confidence, you can read any transformer implementation.

The three things most worth internalising:

1. A tensor is a view into flat memory — shape and strides are a navigation recipe, not the data itself. Most “shape operations” don’t move anything.
2. The @ operator batches over all leading dims and does matmul over the last two. This is how all multi-head attention scoring works.
3. A RuntimeError about shapes or devices always tells you the actual shapes in the message. Read the message before guessing.

Further Reading

PyTorch official documentation

• Tensor tutorial — the official introduction to PyTorch tensors, with interactive examples. Covers creation, indexing, and basic operations: pytorch.org/tutorials/beginner/basics/tensorqs_tutorial.html
• Autograd mechanics — the definitive explanation of how PyTorch’s automatic differentiation works, including the computational graph and gradient accumulation: pytorch.org/docs/stable/notes/autograd.html
• torch.Tensor documentation — the full API reference for every tensor method and attribute: pytorch.org/docs/stable/tensors.html

Understanding memory and performance

• PyTorch internals — Edward Yang — a deep dive into how tensors, storage, and strides work under the hood. Essential reading if you want to understand contiguity and memory layout: blog.ezyang.com/2019/05/pytorch-internals
• PyTorch memory management — the official notes on CUDA memory allocation, caching, and how to track usage: pytorch.org/docs/stable/notes/cuda.html

Transformer-specific tensor patterns

• The Annotated Transformer (Harvard NLP) — a line-by-line walkthrough of the original Transformer paper implemented in PyTorch. Every shape transformation is visible: nlp.seas.harvard.edu/annotated-transformer
• Llama model source (Hugging Face) — reading actual production transformer code is the fastest way to see all these patterns applied. The modeling_llama.py file contains the head-splitting pattern, GQA expansion, RoPE rotation, and causal masking in one place: github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py