PyTorch tensors and operations: Exercises — AI-ML deep dive

A quick vocabulary reminder so nothing below is a surprise:

$rank (T) = n, numel (T) = \prod_{i = 1}^{n} d_{i}$

Here $T$ is a tensor, $n$ is its rank, and $d_{i}$ is the size of axis $i$ . The big $\prod$ symbol just means "multiply all of these together", the same way $\sum$ means "add all of these together".

Level 1 — Recognition

Exercise 1.1 — Read a shape

A tensor is created with torch.zeros(4, 1, 5). State its rank, its shape, and its numel.

Recall Solution

Shape is exactly the arguments: $(4, 1, 5)$ .
Rank = how many numbers in that tuple = $3$ .
numel = multiply them: $4 \times 1 \times 5 = 20$ .

Picture it as 4 sheets, each sheet has 1 row of 5 numbers → $4 \times 1 \times 5 = 20$ little cells.

You want a $(2, 3)$ tensor of samples drawn from the standard normal distribution (mean 0, spread 1). Which of torch.rand(2,3), torch.randn(2,3), torch.ones(2,3) do you pick, and why not the others?

Recall Solution

Pick torch.randn(2,3).

rand gives uniform numbers in $[0, 1)$ — always positive, no bell curve.
ones gives all $1.0$ — no randomness at all.
randn (the extra n = normal) gives the bell-curve numbers we asked for, centred at 0.

Level 2 — Application

Exercise 2.1 — Broadcasting shapes

You add a tensor of shape $(4, 3)$ to a tensor of shape $(3,)$ . Does it work? If so, what is the result shape? Walk the rule.

Recall Solution

Yes, result shape is $(4, 3)$ . Walk the broadcasting rule (align shapes on the right):

Ranks differ ( $2$ vs $1$ ), so pad the smaller with a leading 1: $(3,) \to (1, 3)$ .
Compare axis by axis: $(4, 3)$ vs $(1, 3)$ . Axis-1: $3 = 3$ ✓. Axis-0: $4$ vs $1$ — one of them is 1, so the size-1 side is stretched to 4. ✓
Result: $(4, 3)$ .

See the figure — the row of 3 gets copied down onto every one of the 4 rows.

Exercise 2.2 — Reduction with a chosen axis

Given $x = [142536]$ compute x.sum(dim=0) and x.sum(dim=1). State the resulting shape of each.

Recall Solution

The rule: dim=k means "collapse axis $k$ " — that axis disappears.

dim=0 collapses the rows (the length-2 axis), summing down each column: $[1 + 4, 2 + 5, 3 + 6] = [5, 7, 9], shape (3,) .$
dim=1 collapses the columns (the length-3 axis), summing across each row: $[1 + 2 + 3, 4 + 5 + 6] = [6, 15], shape (2,) .$

Level 3 — Analysis

Exercise 3.1 — Will this matmul run?

For each pair, say whether A @ B is legal and give the output shape: (a) $A : (2, 3)$ , $B : (3, 4)$ (b) $A : (2, 3)$ , $B : (2, 4)$ (c) $A : (10, 2, 3)$ , $B : (10, 3, 4)$ .

Recall Solution

The law of matrix multiply: the inner two dimensions must match; the outer ones become the result.

(a) $(2, \underline{3}) \times (\underline{3}, 4)$ : inner $3 = 3$ ✓ → output $(2, 4)$ . Legal.
(b) $(2, \underline{3}) \times (\underline{2}, 4)$ : inner $3 \neq = 2$ ✗ → illegal, mismatch error.
(c) Leading $10$ is a batch axis (same on both), inner $3 = 3$ ✓ → output $(10, 2, 4)$ . Legal.

Exercise 3.2 — keepdim and broadcasting

A batch x has shape $(4, 3)$ . You want to subtract each column's mean from that column (a mini version of batch norm). Compare x - x.mean(dim=0) versus x - x.mean(dim=0, keepdim=True). Do both work? Why does the second one always work safely?

Recall Solution

x.mean(dim=0) collapses axis 0 → shape $(3,)$ . Then $(4, 3) - (3,)$ broadcasts: $(3,) \to (1, 3) \to (4, 3)$ ✓. Works here.
x.mean(dim=0, keepdim=True) keeps axis 0 as size 1 → shape $(1, 3)$ . Then $(4, 3) - (1, 3)$ broadcasts directly ✓.

Both give the same numbers in this case. But if you'd reduced along dim=1 instead you'd get $(4,)$ , and $(4, 3) - (4,)$ would try to pad to $(1, 4)$ and fail (4≠3). keepdim=True gives $(4, 1)$ , which broadcasts against $(4, 3)$ cleanly. So keepdim=True is the safe habit — the reduced axis stays in place as a size-1 slot ready to stretch.

Level 4 — Synthesis

Exercise 4.1 — Build a linear layer by hand

Input x has shape $(32, 784)$ (32 MNIST samples, each 784 pixels). You want 128 outputs per sample. Design the weight W and bias b shapes, write the forward pass, and give the output shape. Then count how many multiply-adds the matmul does.

Recall Solution

PyTorch convention: $y = xW + b$ with x as $(batch, in)$ .

W must turn 784 → 128, so shape $(784, 128)$ .
b is one number per output, shape $(128,)$ .
Forward: y = x @ W + b.
Shapes: $(32, 784) \times (784, 128) = (32, 128)$ ; then $+ b$ broadcasts $(128,) \to (1, 128) \to (32, 128)$ .
Output shape $(32, 128)$ .

Multiply-adds in the matmul: each of the $32 \times 128$ output entries is a sum over the inner axis of length 784: $32 \times 128 \times 784 = 3211264 multiply-adds.$

Exercise 4.2 — Reshape a batch of images for a dense layer

x = torch.randn(10, 3, 28, 28) (10 RGB images). Flatten everything except the batch axis so you can feed a fully-connected layer. Write the call and give the new shape and its numel (to prove nothing was lost).

Recall Solution

Use x.flatten(start_dim=1) — keep axis 0, merge the rest: $(10, 3, 28, 28) \to (10, 3 \cdot 28 \cdot 28) = (10, 2352) .$ Check numel is conserved: before $= 10 \cdot 3 \cdot 28 \cdot 28 = 23520$ ; after $= 10 \cdot 2352 = 23520$ . ✓ Reshaping only reorganises the same numbers; it never creates or destroys any.

Level 5 — Mastery

Exercise 5.1 — Trace a full mini forward pass

Given the tiny exact tensors $x = [1324], W = [100111], b = [111]$ compute y = x @ W + b fully by hand, then apply ReLU (which replaces every negative entry with 0 — here nothing is negative). Give the final $(2, 3)$ matrix.

Recall Solution

x is $(2, 2)$ , W is $(2, 3)$ → inner $2 = 2$ ✓ → x@W is $(2, 3)$ . Each entry is row-of-x · column-of-W: $x W = [1 \cdot 1 + 2 \cdot 0 3 \cdot 1 + 4 \cdot 0 1 \cdot 0 + 2 \cdot 1 3 \cdot 0 + 4 \cdot 1 1 \cdot 1 + 2 \cdot 1 3 \cdot 1 + 4 \cdot 1] = [132437] .$ Add b = [1,1,1] to every row (broadcast $(3,) \to (2, 3)$ ): $y = [243548] .$ ReLU: all entries already positive, so the output is unchanged: $ReLU (y) = [243548] .$

a = torch.tensor([1.0, 2.0, 3.0])
b = a.view(3, 1)
b[0, 0] = 99.0

What is a afterward, and why? Then state the one-line fix if you wanted a untouched.

Recall Solution

a becomes [99.0, 2.0, 3.0]. view returns a new shape label onto the same underlying memory — it does not copy. Writing through b writes the shared storage, so a sees it too. Fix: break the sharing with a copy — b = a.view(3, 1).clone() (or a.reshape(...).clone()).

Recall Self-check reveals

What does dim=k do in a reduction? ::: Collapses (removes) axis $k$ ; that axis disappears from the shape. Broadcasting aligns shapes on which side? ::: The right; then pads missing leading axes with 1 and stretches size-1 axes. Inner-dimension rule for A @ B? ::: The last axis of A must equal the second-to-last axis of B. Why keepdim=True? ::: It leaves the reduced axis as size 1 so the result broadcasts back against the original. Does view copy data? ::: No — it shares memory; use .clone() to break sharing.

Connections: 3.1.01-neural-network-fundamentals · 3.2.03-backpropagation · 3.3.02-building-modelspytorch · 2.4.01-numpy-fundamentals · 3.4.05-batch-normalization · 3.5.02-convolutional-layers · 4.2.01-gpu-acceleration

Exercises — PyTorch tensors and operations

Level 1 — Recognition

Exercise 1.1 — Read a shape

Exercise 1.2 — Which creator?

Level 2 — Application

Exercise 2.1 — Broadcasting shapes

Exercise 2.2 — Reduction with a chosen axis

Level 3 — Analysis

Exercise 3.1 — Will this matmul run?

Exercise 3.2 — keepdim and broadcasting

Level 4 — Synthesis

Exercise 4.1 — Build a linear layer by hand

Exercise 4.2 — Reshape a batch of images for a dense layer

Level 5 — Mastery

Exercise 5.1 — Trace a full mini forward pass