3.3.1 · D3Deep Learning Frameworks

Worked examples — PyTorch tensors and operations

3,773 words17 min readBack to topic

This is the hands-on drill room for the parent tensor note. The parent told you what the operations are. Here we grind through every kind of situation those operations can create — matching shapes, mismatching shapes, empty tensors, memory traps, gradient tracking — until nothing can surprise you.

Before line one: a tensor is just a box of numbers arranged in a grid. A grid can be a single line (1D), a rectangle (2D), a stack of rectangles (3D), and so on. The word shape is the list of sizes along each direction of that grid. If someone writes shape , picture a rectangle with rows and columns. That mental picture is all you need — every example below leans on it.


The scenario matrix

Here is the full menu of situations tensor operations throw at you. Each later example is tagged with the cell(s) it covers, so by the end every row is checked off.

# Case class What makes it tricky Example
A Matching shapes Plain element-wise add/multiply — the easy baseline Ex 1
B Broadcasting (rank mismatch) One tensor has fewer dimensions than the other Ex 2
C Broadcasting (size-1 stretch) A dimension of size gets stretched to match Ex 2
D Reduction along a chosen axis sum/mean with dim=which axis collapses? Ex 3
E keepdim + broadcast combo Normalisation: reduce, then subtract back Ex 4
F Matrix multiply shape rule Inner dims must match; outer dims survive Ex 5
G Reshape / view / flatten Rearrange without touching the numbers Ex 6
H Degenerate / zero-size inputs Empty tensor, scalar (0D), single element Ex 7
I Memory sharing trap from_numpy, view share memory; clone breaks it Ex 8
J Real-world word problem A full linear layer forward pass Ex 9
K Exam-style twist Broadcast that looks legal but is illegal Ex 10

Example 1 — Case A: matching shapes (the baseline)

Forecast: guess the shape and the numbers before reading on. (Both shapes match, so nothing gets stretched.)

Look at the figure: two rows of three boxes, lined up perfectly. Same shape means we just pair up box-by-box.

Figure — PyTorch tensors and operations
  1. Add position by position. . Why this step? When shapes are identical PyTorch does no stretching — it walks both grids in lockstep. This is the simplest case and the yardstick for everything harder.
  2. Multiply position by position. . Why this step? Element-wise * is not matrix multiplication — it never mixes different positions. This is exactly how activation functions and scaling work.
  3. Output shape stays — element-wise ops never change shape.

Verify: total of the sum is , and indeed . ✓


Example 2 — Cases B + C: broadcasting from a scalar and a row

Forecast: what shape does adding a lone number to a grid produce? What about adding a single row?

The word broadcasting means: when shapes don't match, PyTorch copies the smaller thing to fill the gap, without actually using extra memory. The figure shows the row being copied downward onto every row of .

Figure — PyTorch tensors and operations
  1. Scalar case. has shape (a 0-dimensional tensor — rank , no axes at all). By rule 1 it becomes ; every axis is size , so it legally stretches to cover all boxes. Result rows: and , shape . Why this step? A scalar broadcasts against anything — this is why x * 2 in the parent note "just works".
  2. Row case — align ranks (rule 1). has rank (two axes), has rank (one axis). Prepend a to the smaller rank, so is treated as shape . Why this step? Broadcasting compares shapes from the right. To compare fairly it first pads the missing left-axes of the lower-rank tensor with .
  3. Legality check (rule 2), then stretch (rule 3). vs : rightmost axis (equal, legal); left axis vs — one is , legal, so the size- axis is stretched to . Row is copied onto both rows of . Result rows: and , shape . Why this step? This is precisely the mechanism behind adding a bias vector to a batch — one bias row reused for every sample.

Verify: row-1 sum of (b) is ; from Example 1 that pattern is right. Row-2 sum is . ✓


Example 3 — Case D: reduction — which axis dies?

Forecast: dim=0 sums across... rows or columns? Guess the two answers.

Here is the trap everyone hits: dim=k means "the axis that disappears is axis ". Axis is the vertical (row) direction, axis is the horizontal (column) direction. The figure marks each with a coloured arrow — the arrow points along the axis that gets collapsed.

Figure — PyTorch tensors and operations
  1. dim=0 collapses the vertical axis. We add down each column: . The rows vanish; columns remain, shape . Why this step? Killing axis removes the row-count from the shape. Loss functions do this to average over a batch: "sum away the batch axis."
  2. dim=1 collapses the horizontal axis. We add along each row: . The columns vanish; rows remain, shape . Why this step? Killing axis removes the feature-count. This is how you'd total each sample's features.

Verify: both reductions must reproduce the grand total. and , and the full grid sum is . All three agree. ✓


Example 4 — Case E: keepdim and centering (the normalisation move)

Forecast: the mean over dim=0 has shape . Can you subtract a tensor from a tensor directly? What does keepdim change?

This is the pattern behind batch normalization: reduce along the batch axis, then subtract the result back off every sample.

  1. Column means, dim=0. Column means are , , . Without keepdim this is shape . Why this step? We want each feature's average, and features run across columns, so we collapse the row (batch) axis.
  2. Use keepdim=True so the mean has shape instead of . Why this step? We are about to compute . With shape the mean broadcasts cleanly against (Case C, size-1 stretch). Shape also broadcasts here by luck, but the moment you reduce over dim=1 you'd get , which would not align — keepdim makes the pattern robust every time.
  3. Subtract (broadcasts). Row 1: . Row 2: . Shape . Why this step? After centering, each column's new mean is — exactly what a normalisation layer wants.

Verify: new column means must be : column 1 is . ✓ Sum of the whole centred grid is (means removed all mass): . ✓


Example 5 — Case F: matrix multiplication shape rule

Forecast: which two numbers must match, and which two survive into the output shape?

The rule: for , the inner dimensions must be equal, and the outer dimensions form the output. Write the shapes side by side: . The figure shows the inner s "touching and cancelling", leaving .

Figure — PyTorch tensors and operations
  1. Check the inner dims. 's last axis is ; 's first axis is . They match, so the multiply is legal. Why this step? Each output entry is a dot product over that shared inner axis — if the lengths differed there'd be nothing to pair up.
  2. Output shape = outer dims . Why this step? Row of dotted with column of gives entry ; there are rows and columns, hence .
  3. Compute entry . Row of is ; column of is . Dot product . Why this step? Since is all ones, every entry in row equals the row sum of 's row , which is .

Verify: because is all ones, row of the product should be and row should be (row sums of ). Row-1 sum check: . ✓


Example 6 — Case G: reshape / view / flatten

Forecast: guess the flattened feature count and the value resolves to.

  1. Total elements (this is numel) . Why this step? Reshaping can never create or destroy numbers, so the product of the new shape must equal this. , times , times .
  2. Flatten per image. Keep the batch axis, merge the rest: x.flatten(start_dim=1) gives shape . Why this step? A dense layer wants each sample as one long feature vector, but we must not blend different images together — so we flatten from axis onward, protecting axis .
  3. The -1 placeholder. view(-1, 4) on elements: the known axis is , so becomes , giving . Why this step? means "you do the arithmetic for me." PyTorch fills it so the total () is preserved.

Verify: each new shape's product equals the original element count. . ✓ . ✓


Example 7 — Case H: degenerate & zero-size inputs

Forecast: what's the sum of no numbers? What's the rank of a lone ?

  1. Empty-tensor sum. Shape holds zero elements (numel ). sum returns the additive identity, . Why this step? The sum of an empty set is defined as so that splitting-and-recombining sums stays consistent (an empty sub-batch contributes nothing). Never assume it errors — it quietly returns .
  2. Scalar tensor. torch.tensor(7.0) has shape (empty tuple) and rank — a 0D tensor. Its numel is (one number lives inside, even with no axes). Why this step? A single value still lives in a tensor, but with no axes. This is what a loss value looks like: one number, rank . .item() pulls out the plain Python float .
  3. Squeeze a tensor. .squeeze() removes every axis of size , so , a scalar again holding , rank . Why this step? squeeze deletes only length- axes because those carry no real spread of data; here all three qualify, collapsing to a scalar.

Verify: empty sum ; scalar numel ; squeezed value with rank . All three are checked below. ✓


Example 8 — Case I: the memory-sharing trap

Forecast: does editing a leak into b? Into c?

The word view means a second window onto the same numbers — no copy is made. clone makes a genuine separate copy.

  1. b is a view. It points at the same storage as a. After a[0] = 99, reading b[0] gives . Why this step? view never allocates new memory; it only reinterprets the existing block. Editing through one name is visible through the other — this is the #1 silent bug when people expect a copy.
  2. c is a clone. It owns its own storage, frozen at the instant it was made. c[0] is still . Why this step? clone breaks the link deliberately — use it whenever you need to keep an untouched snapshot.

Verify: b[0] == 99.0 (shared) and c[0] == 1.0 (independent). ✓

import numpy as np, torch arr = np.array([1., 2., 3.]) t = torch.from_numpy(arr) # shares memory with arr arr[0] = 99.0 # edit the NumPy side

now t[0] is ALSO 99.0

safe = torch.from_numpy(arr).clone() # independent snapshot

> **Walk it through.** `torch.from_numpy(arr)` does **not** copy — `t` and `arr` point at one shared block of memory (this is the same link discussed in [[2.4.01-numpy-fundamentals|NumPy fundamentals]]). So editing `arr[0]` mutates `t[0]` underneath you: `t[0]` becomes $99.0$ even though you never touched `t`. The reverse is true too — `t.numpy()` shares memory back. **Fix:** append `.clone()` to break the link the instant you need an independent tensor, as with `safe` above.

> [!mistake] The classic leak
> `t.numpy()` and `torch.from_numpy(arr)` **share memory**. Modify the NumPy side and your tensor mutates underneath you — a bug that hides for hours because the two objects look independent. If you need a truly independent tensor, always break the link explicitly: `torch.from_numpy(arr).clone()`.

---

## Example 9 — Case J: real-world word problem (a full linear layer)

> [!example] Statement
> An MNIST batch has $32$ images, each flattened to $784$ features: input $x$ has shape $(32, 784)$. A linear layer maps $784$ inputs to $128$ outputs, using weight $W$ of shape $(784, 128)$ and bias $b$ of shape $(128,)$. Compute the output shape of $y = xW + b$, and confirm the bias broadcasts correctly. This is the forward pass of [[3.3.02-building-modelspytorch|a PyTorch model layer]] feeding into [[3.2.03-backpropagation|backpropagation]].

**Forecast:** what is $y$'s shape, and how does a $(128,)$ bias reach $32$ different samples?

1. **Matmul shape.** $(32, 784) @ (784, 128)$: inner dims $784=784$ match; outer dims give $(32, 128)$ (Case F, and recall `@` = matmul).
   *Why this step?* Each of the $32$ samples (a row of $x$) is turned into a $128$-dimensional output vector — one row of $y$ per sample.
2. **Bias broadcast.** $b$ has shape $(128,)$, i.e. rank $1$; padded to $(1,128)$, then the size-$1$ axis stretches to $32$ (Cases B+C, and the legality check of Example 2 passes: $128=128$, $32$ vs $1$).
   *Why this step?* Every sample receives the **same** bias vector, added to its $128$ outputs. This is why the layer's convention is $xW$ (not $Wx$): it keeps the batch axis on the left so broadcasting the bias is automatic.
3. **Result** $y$ has shape $(32, 128)$: $32$ samples, each now a $128$-feature vector, ready for the next layer / activation.
   *Why this step?* This exact shape is what feeds [[3.5.02-convolutional-layers|later layers]] or a softmax head; every network is a chain of these. On a GPU this runs in parallel (see [[4.2.01-gpu-acceleration|GPU acceleration]]).

**Verify:** dimension bookkeeping: $(32,784)\times(784,128)\to(32,128)$, and $(32,128)+(1,128)\to(32,128)$. The inner $784$ cancels; batch $32$ and output $128$ survive. ✓

---

## Example 10 — Case K: the exam-style twist (illegal broadcast)

> [!example] Statement
> You try `A + B` with $A$ shape $(3,)$ and $B$ shape $(4,)$. Does it broadcast? Now $A$ shape $(2,3)$ and $B$ shape $(2,)$ — legal or not? Fix the second one so it works.

**Forecast:** decide *legal* or *error* for each before reading.

Broadcasting compares shapes **right-aligned** (the rules boxed in Example 2). Two axes are compatible only if they're **equal** or **one of them is $1$**. No other case is allowed.

1. **$(3,)$ vs $(4,)$.** Rightmost axes: $3$ vs $4$. Not equal, neither is $1$. **Illegal** — PyTorch raises a size-mismatch error.
   *Why this step?* There is no consistent way to stretch a length-$3$ line to length $4$; broadcasting refuses rather than guess.
2. **$(2,3)$ vs $(2,)$.** Right-align: compare $3$ (from $A$) with $2$ (from $B$). $3 \ne 2$ and neither is $1$. **Illegal** — this is the trap: the $2$s look like they "match" but they're in *different* positions after right-alignment.
   *Why this step?* Beginners see a $2$ in both shapes and assume it's fine. But alignment is from the right, so $B$'s $2$ lines up with $A$'s $3$, not $A$'s $2$.
3. **Fix it.** Make $B$ shape $(2, 1)$ (via `B.unsqueeze(1)`). Now right-align: $1$ vs $3$ → stretch (size-1 rule); $2$ vs $2$ → match. **Legal**, output $(2,3)$.
   *Why this step?* Reshaping $B$ to $(2,1)$ puts its $2$ on the correct (row) axis so each row of $A$ gets its own scalar from $B$.

**Verify:** compatibility of the fixed pair: axis-by-axis $(2,3)$ vs $(2,1)$ → $[2{=}2,\ 3\text{ vs }1{\to}\text{stretch}]$ → $(2,3)$. Legal. ✓

---

## Recall

> [!recall]- What is the "rank" of a tensor, and what is `numel`?
> Rank = number of axes (length of the shape tuple); `numel` = total number of elements (product of the shape). Shape $(2,3)$ has rank $2$ and `numel` $6$.

> [!recall]- Which axis vanishes with `sum(dim=1)` on a $(2,3)$ tensor, and what shape remains?
> Axis $1$ (the columns) collapses; shape $(2,)$ remains — one total per row.

> [!recall]- Why does the linear layer use $xW$ with $W$ of shape (in, out) rather than $Wx$?
> So the batch axis stays leftmost: $(\text{batch},\text{in})@(\text{in},\text{out})=(\text{batch},\text{out})$, letting the bias broadcast to every sample automatically.

> [!recall]- $A$ shape $(2,3)$, $B$ shape $(2,)$ — legal to add? Fix?
> Illegal (right-aligned $3$ vs $2$). Reshape $B$ to $(2,1)$; then it broadcasts to $(2,3)$.

> [!recall]- `b = a.view(...)` then `a[0]=99`. What is `b[0]`?
> $99$ — a view shares memory. Use `a.clone()` to get an independent copy.

> [!recall]- What does the infix `@` do in PyTorch, and how is it different from `*`?
> `@` is matrix multiplication (same as `torch.matmul`); `*` is element-wise multiplication. They give different shapes and different numbers.

> [!mnemonic] Broadcast in one line
> **Right-align, then each axis must be Equal or contains a 1 — Else it Errors.** (R-E-1-E.)

Cloze drill:
- The output shape of $(m,k)@(k,n)$ is ==$(m,n)$== — inner $k$ cancels.
- `keepdim=True` keeps the reduced axis as size ==$1$== so the result still broadcasts.
- The sum of an empty tensor is ==$0.0$==, the additive identity.
- The ==rank== of a tensor is its number of axes; `numel` is its total element count.
- The infix `@` means ==matrix multiplication== (equivalent to `torch.matmul`).