PyTorch and Tensor Fluency Crash Course

Doing research in ML/AI, both in theory and in practice, in academia and in industry, will invariably involve using a tensor library at some point. Even just to prototype or test something. Or to better understand how an algorithm works. PyTorch is the tensor library I'm most familiar with and so this tutorial will focus on PyTorch but many concepts apply to other libraries as well (e.g. TensorFlow, JAX).

This is my attempt to document the minimum theory + minimum API knowledge needed to become fluent in tensor manipulations. This is the first blog I've written in which AI has been instrumental in my learning, practicing and documenting my thoughts. I used Claude Fable 5 to specifically write down the PyTorch minimal function list in Section 5 as well as big portions of Section 6. I used Gemini in Antigravity for quick unit-testing and verification of the content.

General tip: The goal of much of this tutorial is to make thinking and reasoning with Tensors easier. That doesn't mean it will be easy. A pen and paper will always help! Also, you haven't mastered Tensor programming until you can, at the very least, reproduce most of the examples in this tutorial from scratch. So practice away!

More helpful resources worth checking out:
1. On shape suffixes: Noam Shazeer's post
2. On broadcasting and Tensor Puzzles: Sasha Rush's GitHub

0. Dimension key (used throughout this doc)

This is the shape-suffix convention from Noam Shazeer's post: pick single letters for logical dimensions, document them once at the top of the file, and end every tensor variable name with its dimension suffix. We'll use this key everywhere below:

B: batch size
T: query sequence length ("time")
M: key/value sequence length (memory; M == T for self-attention)
D: model dimension (d_model)
H: number of attention heads
K: dimension per head (d_model // n_heads, sometimes d_k)
F: feed-forward hidden dimension (usually 4*D)
V: vocabulary size

So x_BTD is a (batch, seq, d_model) tensor and scores_BHTM is (batch, heads, query_len, key_len). The payoff: shape errors become visible in the source code before you ever run it. If you write x_BTD @ w_FD, you can see the inner dims (D vs F) don't line up.

Two mechanical habits that follow from it:

q_BTHK = q_BTD_reshaped          # after splitting heads
q_BHTK = q_BTHK.transpose(1, 2)  # the suffix tells you what transpose did

Every time a tensor's shape changes, its name changes. That's the whole convention.

1. The one mental model: a tensor = memory + shape + strides

A tensor is a flat 1-D block of memory plus two tuples of metadata:

shape: how big each dimension is
stride: how many elements you skip in flat memory to move one step along each dimension

x = torch.arange(12).reshape(3, 4)
x.shape    # (3, 4)
x.stride() # (4, 1)  -> element [i, j] lives at flat index i*4 + j*1

For another example, a $(2,3,4,5)$ tensor has stride $(60, 20, 5, 1)$. For dimension $i$ (0-indexed) and size $s_i$, the stride is $\prod_{j=i+1}^{n-1} s_j$.

Almost every "shape function" just rewrites this metadata without touching memory. These are views (free, O(1), share storage):

reshape / view (when contiguous): you can reshape a tensor to any shape, provided the new strides are compatible with the old ones.
transpose, permute: swap two dimensions.
indexing with None: adds a singleton dimension. For example: x.shape -> (B,T,D) can be reshaped to (B,T,1,D) using x[:, :, None, :].
slicing: preserves strides but reduces the size of the sliced dimension. For example: x[:, :T//2, :].
expand: stretches a singleton dimension to match another tensor's shape. For example: x.shape -> (B,T,1) can be expanded to (B,T,M) using x.expand(-1, -1, M).

Again, we just rewrite the metadata (shape and strides). We don't touch the memory. Using this viewpoint we can understand the classic gotchas:

y = x.transpose(0, 1)   # shape (4, 3), stride (1, 4) — same memory, reordered strides
y.view(12)              # ERROR: strides no longer describe a contiguous walk
y.contiguous().view(12) # fine: .contiguous() copies memory into standard order
y.reshape(12)           # fine: reshape copies for you when it must

And why broadcasting is free: expand sets a dimension's stride to 0, so the same memory element is read repeatedly with no copy. (repeat, by contrast, actually copies.)

This is the conceptual question behind "what's the difference between view and reshape?", "why is expand free but repeat isn't?", and "why did my .view throw an error after a transpose?" One model answers all three.

2. The one rule: broadcasting

This is a super valuable rule to remember. I'm borrowing the phrasing from Sasha Rush's GitHub.

When two tensors of different shapes meet in an elementwise op, PyTorch:

Right-aligns the shapes.
Left-pads the shorter one with 1s.
For each dimension: sizes must be equal, or one of them must be 1. A size-1 dim is stretched (stride 0) to match.

(B, T, D)  *  (D,)      -> (B, T, D)    # per-channel scale
(B, T, 1)  +  (B, 1, M) -> (B, T, M)    # pairwise combination
(T, 1)     >= (1, M)    -> (T, M)       # comparison table

The workhorse idiom is inserting size-1 dims with None (identical to unsqueeze):

a_I, b_J = torch.arange(3), torch.arange(4)
outer_IJ = a_I[:, None] * b_J[None, :]   # (3,1) * (1,4) -> (3,4)

This one idiom — arange + [:, None] + a comparison — generates every mask you will ever need:

i_T = torch.arange(T)
causal_TT = i_T[:, None] >= i_T[None, :]          # lower-triangular True
seq_mask_BT = i_T[None, :] < lengths_B[:, None]    # padding mask from lengths

Read i[:, None] >= j[None, :] as: "build the full table of (query index i, key index j) pairs and ask, for each pair, may query i look at key j?" Broadcasting is how you replace a double for-loop with a table.

To practice this, try solving Sasha Rush's Tensor Puzzles. The rules there only allow arange, where, @, arithmetic, comparison, .shape, and indexing — because broadcasting over those primitives is sufficient to rebuild sum, cumsum, eye, bincount, and the rest. More on that in section 6.

3. The one notation: einsum

einsum notation makes reasoning with high-dimensional tensors easier. In deep learning we often deal with 4 dimensional tensors and have to do various contractions, transpositions and batch multiplications between them. These are often easy to understand on an individual level with 2D matrices, but when mini-batches and training data come into play, it can get confusing.

Once the mathematical intent is clear to you, einsum can take over and take care of all the bookkeeping of transpositions and batching.

The notation is not very complicated: you name each dimension of each input with a letter, then declare the output's letters.

A letter that appears in the inputs but not in the output is summed over (contracted).
A letter repeated across inputs means those dims are multiplied elementwise (then summed if absent from output).
Letters in the output determine its shape and axis order — so einsum is also a transpose.

Imagine the classic einstein sum notation, just without the sums. For example, for matrix multiplication we would write
$$ C_{ij} = \sum_k A_{ik}B_{kj} \to A_{ik}B_{kj} $$
Here are some additional examples:

torch.einsum('ij,jk->ik',  A, B)      # matmul: j is contracted
torch.einsum('btd,dk->btk', x, W)     # batched linear layer
torch.einsum('bthk,bmhk->bhtm', q, k) # attention scores (contract k, keep heads)
torch.einsum('bhtm,bmhk->bthk', p, v) # weighted sum of values
torch.einsum('ij->ji',    A)          # transpose
torch.einsum('ii->',      A)          # trace
torch.einsum('i,j->ij',   a, b)       # outer product
torch.einsum('bd,bd->b',  x, y)       # per-row dot product

Why memorize this instead of ten specialized functions? Because einsum replaces matmul, bmm, mm, mv, dot, outer, tensordot, and half your transpose calls, and it pairs perfectly with shape suffixes — the einsum string is the shape suffix of each argument. When you write attention with suffix-named variables, the einsum strings basically write themselves:

scores_BHTM = torch.einsum('BTHK,BMHK->BHTM', q_BTHK, k_BMHK) / K**0.5

Rule of thumb: use @ for plain matmuls of 2-D (or trivially batched) tensors; reach for einsum the moment more than one dimension needs to survive a contraction (heads!) or you'd otherwise need transpose gymnastics.

4. The minimal function set (~25 things)

Here are most of the functions one needs when working with PyTorch. More contrived examples are usually unnecessary or can be looked up. I won't go in great depth here as this tutorial is supposed to be more of a cheat sheet, but make sure to go through the documentation or ask AI for explanations if you're not familiar with some of the functions below.

Creation (4)

torch.arange(n)              # 0..n-1 — THE loop replacement
torch.tensor([1, 2, 3])      # literal
torch.zeros(s); torch.ones(s)
torch.randn(s)               # for weight init: torch.randn(D, F) * 0.02

Shape (5)

x.reshape(B, T, H, K)        # split/merge dims (use over .view; copies only if needed)
x.transpose(1, 2)            # swap two dims;  x.permute(...) for arbitrary reorder
x[:, None]                   # insert size-1 dim (== unsqueeze); x.squeeze(d) removes
x.expand(B, T, M)            # free broadcast of size-1 dims (no copy). Beware! If you modify the original, then the expanded view also changes! Broadcasting (see above) is preferred typically.
torch.cat([a, b], dim=-1); torch.stack([a, b])   # cat joins existing dim, stack makes new one

Elementwise math (4 groups)

+ - * / **                   # all broadcast
torch.exp, torch.log, torch.sqrt, torch.tanh
torch.maximum(a, b)          # elementwise max of two tensors (flash attention needs this)
==  <  <=  >  >=             # comparisons -> bool tensors (mask factories)

Reductions (4) — always think about `dim=` and `keepdim=`

x.sum(dim=-1, keepdim=True)      # keepdim=True keeps a size-1 dim so you can broadcast back
x.max(dim=-1, keepdim=True)      # returns (values, indices)! use .values
x.mean(dim=-1); x.argmax(dim=-1)
x.cumsum(dim=-1)                 # cumulative sum: [a, b, c] -> [a, a+b, a+b+c]

keepdim=True is not optional trivia — it is what makes x - x.max(-1, keepdim=True).values broadcast correctly. Forgetting it (or forgetting .values on max) is the #1 live-coding bug.

Selection & masking (4)

torch.where(cond, a, b)               # vectorized if-statement
scores.masked_fill(~mask, -torch.inf) # the canonical causal-mask move
x[mask]                               # boolean indexing (selects, flattens)
emb_VD[ids_BT]                        # integer "fancy" indexing -> (B, T, D): embedding lookup IS indexing
                                      # e.g. x = torch.tensor([10, 20, 30]); x[[2, 0, 0]] -> [30, 10, 10]

Linear algebra (2)

a @ b            # matmul; batches over leading dims
torch.einsum(...)

Autograd & training (6)

w = torch.randn(D, F, requires_grad=True)   # or nn.Parameter inside nn.Module
loss.backward()                             # populates .grad on every leaf
with torch.no_grad(): ...                   # for the update step / eval
F.softmax(x, dim=-1); F.cross_entropy(logits_ND, targets_N)   # note: cross_entropy takes raw logits
opt = torch.optim.SGD(model.parameters(), lr=1e-3)
opt.zero_grad(); loss.backward(); opt.step()   # the sacred order

And the two nn building blocks worth knowing cold: nn.Linear(D, F) (stores weight as (F, D), computes x @ W.T + b) and nn.Embedding(V, D) (a learnable lookup table = fancy indexing into a (V, D) matrix).

That's it. Notice what's not on the list: no scatter, no unfold, no einops (nice, but einsum + reshape covers it), no bmm, no index_select. All expressible from the above.

5. Five idioms that generate everything

These are compositions of the primitives above. Internalize them as moves, like chess tactics.

Idiom 1 — arange + comparison = any mask

i_T = torch.arange(T)
causal_TT   = i_T[:, None] >= i_T[None, :]           # query i may see key j iff j <= i
pad_mask_BM = torch.arange(M)[None, :] < lens_B[:, None]
window_TT   = (i_T[:, None] - i_T[None, :]).abs() <= w  # local/sliding-window attention

Idiom 2 — comparison + matmul = one-hot, and one-hot = sum/scatter/bincount

A one-hot matrix is just a comparison table:

onehot_NV = (ids_N[:, None] == torch.arange(V)[None, :]).float()

Now watch what matmul with it does:

counts_V = onehot_NV.sum(0)          # bincount
sums_V   = onehot_NV.T @ values_N    # scatter_add: values routed to their bins. Write this to see it if it's not obvious.
picked_ND = onehot_NV @ emb_VD       # embedding lookup as a matmul.

Idiom 3 — subtract-max, exp, normalize = stable softmax

A standard trick is to subtract the maximum before exponentiating to prevent overflow.

def softmax(x, dim=-1):
    z = x - x.max(dim=dim, keepdim=True).values   # doesn't change the answer, prevents overflow
    e = torch.exp(z)
    return e / e.sum(dim=dim, keepdim=True)

Idiom 4 — matmul with a triangular matrix = cumulative ops

tri_TT = (torch.arange(T)[:, None] >= torch.arange(T)[None, :]).float()
cumsum_T = tri_TT @ x_T

Prefix sums, running means, "attend to everything before me" — all are matmuls against a mask.

Idiom 5 — reshape + transpose = split/merge heads

q_BTHK = q_BTD.reshape(B, T, H, K)      # split: D -> (H, K), because D == H*K
q_BHTK = q_BTHK.transpose(1, 2)         # move H next to B so matmul batches over (B, H)
out_BTD = out_BHTK.transpose(1, 2).reshape(B, T, D)   # merge back (contiguous handled by reshape)

Say the words "reshape splits, transpose moves, reshape merges" and multi-head attention stops being scary.

Note: we cannot do this with einsum because it cannot split or merge dimensions.

6. Case Studies

The principles above show up everywhere in modern deep learning. Here are a few canonical examples below. When thinking about these or writing this code from scratch I first start from the core line (e.g. the matmul in scores_BHTM = q @ k.T for attention) and then think about how to get there.

6a. Attention — self, causal, cross (one function, three configs)

The three variants are the same code; only the source of K/V and the mask differ:

self-attention: q, k, v all from x, no mask
causal self-attention: same, plus the triangular mask
cross-attention: q from decoder x_BTD, k and v from encoder mem_BMD, usually no causal mask

import torch, torch.nn.functional as F

def attention(x_BTD, ctx_BMD, w_q_DD, w_k_DD, w_v_DD, w_o_DD, H, causal=False):
    B, T, D = x_BTD.shape
    M = ctx_BMD.shape[1]
    K = D // H

    q_BHTK = (x_BTD   @ w_q_DD).reshape(B, T, H, K).transpose(1, 2)
    k_BHMK = (ctx_BMD @ w_k_DD).reshape(B, M, H, K).transpose(1, 2)
    v_BHMK = (ctx_BMD @ w_v_DD).reshape(B, M, H, K).transpose(1, 2)

    # Each head first computes the QK^T product individually...
    scores_BHTM = q_BHTK @ k_BHMK.transpose(-2, -1) / K**0.5 # can you do it via einsum?
    if causal:
        mask_TM = torch.arange(T)[:, None] >= torch.arange(M)[None, :]
        scores_BHTM = scores_BHTM.masked_fill(~mask_TM, -torch.inf)
    p_BHTM = F.softmax(scores_BHTM, dim=-1)

    out_BHTK = p_BHTM @ v_BHMK
    out_BTD = out_BHTK.transpose(1, 2).reshape(B, T, D)
    return out_BTD @ w_o_DD

# self:  attention(x, x, ...)            cross: attention(x, mem, ...)
# causal self: attention(x, x, ..., causal=True)

Why / sqrt(K) (keeps score variance ~1 so softmax isn't saturated at init); why mask with -inf before softmax and not zero out probabilities after (rows must still sum to 1 over allowed positions); why softmax over dim=-1 (each query normalizes over keys). Verify your implementation against F.scaled_dot_product_attention.

6b. MLP forward + backward by hand — the four identities

Manual backprop typically reduces to four memorized derivatives plus the chain rule. For more complex situations you might have to do the derivations yourself (and write a small helper script to check that you did them correctly). Given upstream gradient dY (same shape as Y):

1. Linear      Y = X @ W        ->  dX = dY @ W.T        dW = X.T @ dY
2. Add bias    Y = X + b        ->  dX = dY              db = dY.sum(over broadcast dims)
3. ReLU        Y = max(X, 0)    ->  dX = dY * (X > 0)
4. Softmax+CE  L = CE(softmax(Z), y)  ->  dZ = (P - onehot(y)) / N

Sanity mnemonics: shapes force the formulas — dW must have W's shape (D, F), and the only way to build (D, F) from X:(N, D) and dY:(N, F) is X.T @ dY. Rule 2 generalizes to all broadcasting: the gradient sums over every dimension that was broadcast (broadcast forward ⇒ sum backward, and vice versa).

The derivation of rule 4 is a bit annoying but really just relies on the quotient rule to find the gradient of a softmax: $\frac{\partial p_k}{\partial z_j} = p_k(\delta_{kj} - p_j)$.

Full 2-layer MLP with manual backward:

# forward
h1_NF = x_ND @ w1_DF + b1_F
a_NF  = torch.maximum(h1_NF, torch.tensor(0.0))
z_NC  = a_NF @ w2_FC + b2_C
p_NC  = softmax(z_NC)                                  # from Idiom 3
loss  = -torch.log(p_NC[torch.arange(N), y_N]).mean()  # fancy indexing picks true-class probs

# backward (chain rule, bottom-up)
dz_NC  = (p_NC - onehot(y_N, C)) / N     # identity 4
dw2_FC = a_NF.T @ dz_NC                  # identity 1
db2_C  = dz_NC.sum(0)                    # identity 2
da_NF  = dz_NC @ w2_FC.T
dh1_NF = da_NF * (h1_NF > 0)             # identity 3
dw1_DF = x_ND.T @ dh1_NF
db1_F  = dh1_NF.sum(0)
dx_ND  = dh1_NF @ w1_DF.T

Always verify against autograd

loss.backward()
assert torch.allclose(dw1_DF, w1_DF.grad, atol=1e-5)

6c. Attention backward — one new identity

Attention backward = identity 1 applied to every matmul, plus the softmax backward (for a softmax that isn't fused with cross-entropy):

P = softmax(S, dim=-1)
dS = P * (dP - (dP * P).sum(dim=-1, keepdim=True))

Derivation sketch you should be able to reproduce: dS_i = Σ_j dP_j · ∂P_j/∂S_i, with ∂P_j/∂S_i = P_j(δ_ij − P_i); collect terms and the row-wise sum appears. Then the rest of attention backward is mechanical:

# forward pieces:  S = Q @ K.T / sqrt(d);  P = softmax(S);  O = P @ V
dV = P.transpose(-2, -1) @ dO            # identity 1
dP = dO @ V.transpose(-2, -1)
dS = P * (dP - (dP * P).sum(-1, keepdim=True))
dQ = dS @ K / d**0.5
dK = dS.transpose(-2, -1) @ Q / d**0.5

(Masked positions take care of themselves: P is 0 there, so dS is 0 there.) Verify with torch.autograd.grad.

6d. Flash attention — the online softmax

Flash attention is is Idiom 3 computed incrementally over key blocks so the full T×M score matrix never materializes. Per query, sweep over blocks of keys keeping three running stats — max m, normalizer l, un-normalized output accumulator acc:

def flash_attention(q_TK, k_MK, v_MK, block=64):        # single head, single query block
    T, Kd = q_TK.shape
    m_T   = torch.full((T, 1), -torch.inf)
    l_T   = torch.zeros(T, 1)
    acc_TK = torch.zeros(T, Kd)
    for j in range(0, k_MK.shape[0], block):
        s_TB   = q_TK @ k_MK[j:j+block].T / Kd**0.5     # scores vs this key block only
        m_new  = torch.maximum(m_T, s_TB.max(-1, keepdim=True).values)
        scale  = torch.exp(m_T - m_new)                 # correct old stats to the new max
        p_TB   = torch.exp(s_TB - m_new)
        l_T    = l_T * scale + p_TB.sum(-1, keepdim=True)
        acc_TK = acc_TK * scale + p_TB @ v_MK[j:j+block]
        m_T    = m_new
    return acc_TK / l_T                                  # normalize once at the end

The key insight: exp(s − m_new) = exp(s − m_old) · exp(m_old − m_new), so a change of running max is fixable by multiplying previously accumulated quantities by exp(m_old − m_new). Memory drops from O(T·M) to O(block), which is the point — attention is memory-bandwidth-bound, and this keeps everything in SRAM. Verify against your section-6a implementation.

6e. Training loop with SGD

model = MyTransformer(...)                                # nn.Module with __init__ + forward
opt = torch.optim.SGD(model.parameters(), lr=1e-2)

for step, (x_BT, y_BT) in enumerate(loader):
    logits_BTV = model(x_BT)
    loss = F.cross_entropy(logits_BTV.reshape(-1, V), y_BT.reshape(-1))
    opt.zero_grad()
    loss.backward()
    opt.step()

Know why each line exists: zero_grad because .backward() accumulates into .grad rather than overwriting; the reshape because cross_entropy wants (N, V) logits and (N,) targets; cross_entropy takes raw logits (it fuses log_softmax + NLL for numerical stability — never softmax first). And be able to do it with no optimizer at all:

with torch.no_grad():
    for p in model.parameters():
        p -= lr * p.grad
        p.grad = None

The no_grad is needed so the update itself isn't recorded into the autograd graph.