Python & NumPy
Array programming, broadcasting, and vectorized computation — taught from zero, starting with why an ordinary Python list is the wrong shape for data, and built up to the array thinking every ML framework is written in.
Start here — what this is really about
Every model you will ever train is, underneath, a pile of numbers arranged in grids, being multiplied and added billions of times. An image is a grid of brightness values. A batch of training examples is a grid of rows. A layer's weights are a grid. Training is arithmetic on grids.
Python — the language everyone writes ML in — is genuinely bad at doing arithmetic on millions of numbers one at a time. It is slow, by a factor of hundreds. So the entire scientific-Python world is built on one library that fixes exactly this: .
NumPy gives you one new kind of object — an array, a grid of numbers all of the same type, packed side by side in memory. Then it gives you one new habit — stop writing loops; describe what you want done to the whole grid at once, and let compiled C code run the loop for you at full speed.
That second habit has a name — vectorization — and it is the single most important skill on this page. Everything else (shapes, dtypes, broadcasting, views) exists to make it work.
The page adapts to you. By default it teaches from first principles, with no notation until you already understand the idea it stands for. Flip the Depth switch at the top to reveal the formal definitions, the memory model, and the edge cases. Nothing is hidden for good — the deeper material sits behind the Go deeper panels so you can open it the moment you're curious. Every code cell runs right here in your browser, so change the numbers and see what breaks.
PyTorch tensors, TensorFlow tensors, JAX arrays and pandas DataFrames all copy NumPy's rules — the same shapes, the same broadcasting, the same indexing. Learn them once here and you have already learned them for every framework you will touch later. NumPy is also the library the labs assume you can read fluently in an interview.
Step 1 — Why an ordinary Python list is the wrong container
You already know Python lists: [1, 2, 3]. They are wonderfully flexible — a list can hold a number, a string, another list, all at once. That flexibility is exactly what makes them slow.
A Python is a shopping bag: it holds anything, in any order, and each item is really just a note saying where the item actually lives. To add up the contents, you must follow every note to a different shelf. A NumPy is an egg carton: fixed number of identical slots, all touching, all the same size. To add up the contents you walk down one straight line of memory. Same eggs — wildly different walking.
That memory picture explains the two differences you will feel immediately:
- Lists are slow for math, because every single element is a separate Python object that must be located, unboxed, added, and re-boxed.
- Lists don't do math at all in the way you would hope.
[1, 2, 3] * 2does not double your numbers — it gives you a longer list,[1, 2, 3, 1, 2, 3].
Say you want to double every number in [1, 2, 3].
The list way — you must spell out the loop, one element at a time:
- Start an empty result list.
- Take
1, multiply by 2, get2, append it. - Take
2, multiply by 2, get4, append it. - Take
3, multiply by 2, get6, append it. - Result:
[2, 4, 6]. Three separate trips, all steered by Python.
The array way — you state the operation once, for the whole grid:
np.array([1, 2, 3]) * 2 → array([2, 4, 6])
Same answer. But you never wrote a loop, and the loop that did run happened inside compiled C, roughly a hundred times faster per element.
Run both and watch the difference in behaviour — this is the first NumPy surprise for everyone coming from plain Python:
Why does multiplying a Python list by 2 repeat it, while multiplying a NumPy array by 2 doubles its values?
Hint: Think about what each container was designed to hold.
Step 2 — The three questions to ask any array
Whenever you meet an array — yours, a colleague's, one that just crashed — ask three questions in this order. Nearly every NumPy bug is answered by one of them.
- What shape is it? How big is the grid, along each direction?
- What dtype is it? What kind of number is in each slot?
- Which axis am I operating along? Which direction of the grid do I mean?
Shape — the size of the grid
The is a tuple of sizes, read outermost-first. The number of entries in that tuple is the array's .
(5,)— a flat list of 5 numbers. One dimension. Note the lonely comma: it's Python's way of writing a one-element tuple, and its presence is how you tell a 1-D array of 5 from something else.(3, 4)— a table with 3 rows and 4 columns. Two dimensions, 12 numbers total. Row count first, always.(256, 256, 3)— a colour image: 256 pixels tall, 256 wide, and 3 numbers per pixel (red, green, blue). 196,608 numbers.(32, 3, 224, 224)— one batch of training data, in the layout PyTorch expects: 32 images, 3 colour channels each, 224 pixels tall, 224 wide. That single array holds 4.8 million numbers, and a training step operates on all of them in one go.
The habit to build: whenever code misbehaves, print the shapes before you print the values. Nine times in ten, the bug is a shape you assumed and did not check.
An image is the friendliest example, because you can see the grid. Below is a real 8×8 array of brightness values, drawn as a picture: every cell is one number, and every number is one pixel. Hover to read the value stored in each cell.
Dtype — what kind of number sits in each slot
Every element of an array has the same type, and that type is fixed when the array is made. That's the .
The ones that matter in practice:
| dtype | Bytes each | What it's for |
|---|---|---|
float64 | 8 | NumPy's default for decimals. Maximum precision, maximum memory. |
float32 | 4 | The deep-learning standard. Half the memory, plenty of precision for gradients. |
float16 / bfloat16 | 2 | Mixed-precision training on GPUs — fast, and delicate. |
int64 | 8 | Default for whole numbers on Linux and macOS (Windows historically defaulted to 32-bit). |
bool | 1 | Masks — the True/False arrays you filter with. |
Multiply by a billion. A model with one billion parameters stored in float64 needs 8 GB just to hold the weights — before gradients, before optimizer state, before activations. The same model in float32 needs 4 GB, and in float16 only 2 GB. Choosing a dtype is choosing whether your model fits on the GPU at all, which is why every deep-learning framework defaults to float32 while NumPy defaults to float64.
Now inspect some real arrays — shape, dtype, dimensions and the memory they occupy:
You have an array of shape (32, 3, 224, 224) representing a batch of images. How many numbers does it hold, and what does the 3 most likely mean?
Step 3 — Vectorization: say it once, for the whole array
Here is the habit that separates people who use NumPy from people who merely have it installed.
A loop says: take element 0, multiply it, store it; take element 1, multiply it, store it; … — a million tiny instructions, each one interpreted by Python. says: multiply these two million-element arrays. One instruction. Python hands the whole job to compiled C, which runs the loop with no interpreter overhead, no type checks, and often several numbers at a time using the CPU's SIMD units.
The loop is walking to the counter, ordering one sandwich, carrying it back, and repeating a hundred times — the walking dominates, not the sandwich-making. Vectorization is phoning the deli once and saying a hundred of the usual. The kitchen is the same speed; you have removed a hundred round trips. In NumPy, those round trips are Python's interpreter overhead, and they are typically 90-plus percent of the time.
The operations that do this are called , and you have already used several: +, *, np.exp, np.sqrt, np.maximum.
Task: given lists of predicted prices and actual prices, compute the squared error for each — the raw material of a loss function.
Loop version, thinking one element at a time:
- Make an empty results list.
- For each index
i: takepred[i], subtractactual[i], square it, append it. - Repeat for all N examples.
Vectorized version, thinking in whole arrays:
diff = pred - actual— one subtraction, done to every element at once.sq = diff ** 2— one squaring, done to every element at once.mse = sq.mean()— one reduction over the lot.
Same three lines of thought, but each now applies to the entire dataset rather than to one element, and the loop happens in C. For a million examples the second version is measured below at roughly 200 times faster — and it is also shorter and, once your eye adjusts, easier to read.
Measure it yourself, right now, in your own browser. The absolute numbers will differ from the chart above — this cell runs Python compiled to WebAssembly — but the shape of the answer will not:
Your training loop iterates over 50,000 examples in Python and applies the same three arithmetic operations to each. Roughly what will you gain by rewriting it as array operations, and where does the gain come from?
Hint: Count what Python has to do per element versus per array.
Step 4 — Broadcasting: the rule that deletes the rest of your loops
Elementwise operations were defined for arrays of the same shape. But real code constantly combines different shapes: subtract one mean from every row; add one bias vector to every example in a batch; scale a whole image by a single number. is the rule that makes those work without loops and without copies.
When shapes disagree, NumPy asks a simple question along each axis: are these the same size, or is one of them 1? If one of them is 1, that axis gets stretched — the single value is reused for every position along it. And the stretch is a polite fiction: NumPy sets that axis's stride to zero and re-reads the same memory, so a value broadcast across a million rows still occupies one slot.
You have one stamp and a page of a hundred boxes. You do not carve a hundred stamps — you press the same one into every box. Broadcasting is that: one row of biases pressed into every row of the batch, one mean pressed down every column. The stamp is never duplicated, just reused.
Add a column of shape (3, 1) to a row of shape (4,):
col = [[0], [10], [20]] and row = [0, 1, 2, 3].
Step 1 — line the shapes up from the right. The shorter shape is padded with 1s on its left:
| axis 0 | axis 1 | |
|---|---|---|
col | 3 | 1 |
row (padded) | 1 | 4 |
| result | 3 | 4 |
Step 2 — check each axis. Axis 0: sizes 3 and 1 — one of them is 1, so it stretches to 3. Axis 1: sizes 1 and 4 — the 1 stretches to 4. Both axes are compatible, so the answer has shape (3, 4).
Step 3 — fill in. Every result cell is its row's value from col plus its column's value from row:
| 0 | 1 | 2 | 3 | |
|---|---|---|---|---|
| 0 | 0 | 1 | 2 | 3 |
| 10 | 10 | 11 | 12 | 13 |
| 20 | 20 | 21 | 22 | 23 |
Twelve additions, one line of code, zero loops, and no copy of either input was ever made.
That table, drawn:
The whole rule fits in two lines, and it is worth memorising:
- Compare shapes from the rightmost axis leftwards, padding the shorter shape with 1s on the left.
- Two axes are compatible if they are equal, or if one of them is 1. Anything else is an error.
Array A has shape (5, 1, 6) and array B has shape (3, 6). What is the shape of A + B?
Now use broadcasting for the job it does most often in ML — centring a dataset so every feature has mean zero:
Step 5 — Indexing and slicing, and the view-versus-copy trap
Getting numbers out of an array uses the same square brackets as lists, with one index per axis, separated by commas.
The basics, on a 2-D array A:
A[0, 2]— a single element: row 0, column 2.A[1]— an entire row (shape drops by one dimension).A[:, 0]— an entire column. The bare colon means everything along this axis.A[0:3, 2:5]— a rectangular block: rows 0 up to but not including 3, columns 2 up to but not including 5.A[A > 5]— every element greater than 5, as a flat array. This is a , and it is how you filter data without a loop.
You have a grayscale image img of shape (256, 256) and want the 64×64 patch in the top-left corner, then the same patch's average brightness.
- Rows first, columns second — the convention matches the shape tuple:
patch = img[0:64, 0:64]. - Check the shape before trusting it:
patch.shapeis(64, 64). If it came out(64,)you indexed with one index instead of two. - Average it with a reduction over both axes at once:
patch.mean()— a single number. - Want the bottom-right corner instead? Negative indices count from the end:
img[-64:, -64:].
No pixel was copied in step 1 — and that fact is about to matter a great deal.
patch = img[0:64, 0:64] does not duplicate those pixels. It hands you a — a second set of shape and stride numbers pointing at the same bytes. So patch[0, 0] = 255 also changes img[0, 0]. This design is a gift (slicing a 4 GB array costs nothing) and a trap (you can corrupt your training data by normalizing a slice in place). When you want independence, say so: patch = img[0:64, 0:64].copy().
You slice a batch of training images with batch[0:8] and then divide that slice by 255 using the in-place operator. What happened to the original batch, and why?
Hint: Did the slice own its memory?
Step 6 — Reductions and the axis argument
A takes many numbers and returns fewer: sum, mean, max, min, std, argmax. The only hard part is the axis argument, and there is one sentence that makes it click.
axis=0 does not mean along the rows, whatever your intuition insists. It means collapse axis 0 — squash the array down that direction until it is gone. Start from a (4, 3) array: summing with axis=0 removes the 4 and leaves shape (3,) — one number per column. Summing with axis=1 removes the 3 and leaves (4,) — one number per row. Ask which number in the shape tuple do I want to eliminate, and you will never get it backwards again.
Take a (2, 3) array — 2 examples, 3 features:
| f0 | f1 | f2 | |
|---|---|---|---|
| example 0 | 1 | 2 | 3 |
| example 1 | 4 | 5 | 6 |
A.sum(axis=0)— kill axis 0 (the examples). Add down each column: , , . Result[5, 7, 9], shape(3,)— one total per feature.A.sum(axis=1)— kill axis 1 (the features). Add across each row: , . Result[6, 15], shape(2,)— one total per example.A.sum()— no axis named, so kill them all: , a single number of shape().
Sanity check that costs nothing: both of the first two results add up to 21.
You have activations of shape (32, 10) — a batch of 32 examples, 10 class scores each — and you want the predicted class for every example. Which axis do you take the argmax over, and what shape comes back?
Hint: Which number in the shape tuple should be gone from the answer?
X has shape (1000, 20). What shape does X.mean(axis=0, keepdims=True) return, and why would you want keepdims?
Step 7 — Putting it together: a neural network layer with no loops
Everything so far now pays off in one page of code. A layer of a neural network takes a batch of inputs, multiplies by a weight matrix, adds a bias, and squashes the result into probabilities. Written with loops it is forty lines; written with arrays it is four.
Where this shows up everywhere else
Each of these is something on this page, wearing different clothes:
- A PyTorch tensor is a NumPy array with a device and a gradient attached.
.shape,.reshape, broadcasting andaxis(renameddim) all behave the same way. - A pandas DataFrame is a set of NumPy arrays with labels bolted on;
df.valueshands you the array back. That's the next stop, Pandas & DataFrames. - Image preprocessing — cropping, flipping, normalizing — is slicing, striding, and broadcasting.
- A whole training batch exists because one matrix multiply over examples beats matrix multiplies, exactly as measured in Step 3.
- Every interview screen for an ML role asks you to vectorize something, or to explain why the shapes in a snippet do not line up.
Explain to a friend who knows basic Python why NumPy is fast, using the egg-carton and lunch-order pictures rather than symbols. Then explain broadcasting to them by describing what happens when a (4, 3) array meets a (3,) array. If you stall on either, that is exactly the section to reread.
- A NumPy array is a grid of numbers of one dtype, packed contiguously in memory — unlike a Python list, which is a bag of pointers to separate objects.
- Ask three questions about any array: shape (how big along each axis), dtype (what kind of number, and therefore how much memory), and which axis you are operating along.
- Vectorization means describing an operation on whole arrays instead of writing a loop. The speedup — around 200× in the measured example — comes from removing per-element Python overhead, not from faster arithmetic.
- Broadcasting combines different shapes by comparing axes from the right and stretching any axis of size 1. It is what makes adding a bias to a batch, or centring a dataset, a single line with no copies.
- Slicing gives a view that shares memory; fancy indexing and
.copy()give a copy. Writing through a view edits the original — a gift and a trap. - The
axisargument names the dimension that disappears;keepdims=Trueleaves a 1 behind so the result broadcasts back. - A neural network layer is
X @ W + bfollowed by a stable softmax — every idea above, in four lines and zero loops.
Practice — and how to make it stick
Three research-backed habits, built into this platform:
• Retrieval practice: attempt each exercise below before scrolling back up — pulling shapes and rules out of memory builds far more durable knowledge than re-reading them.
• Spaced repetition: mark this topic complete and it joins your Review queue, resurfacing right before you would otherwise forget it.
• Interleaving: mix these with problems from Linear Algebra rather than grinding array syntax in one block — messier practice, sturdier memory. Matrix shapes and NumPy shapes are the same skill seen twice.
- Predict, then check: for each pair of shapes —
(3, 4)with(4,);(3, 4)with(3,);(3, 1)with(1, 4);(2, 3, 4)with(3, 1)— write down whether they broadcast and what shape results. Then verify withnp.broadcast_shapes. Every one you get wrong is a bug you would have shipped. - Vectorize it: write a loop that computes the Euclidean distance from one point to each of 1,000 others, then rewrite it with broadcasting and time both. Then do the harder version — all pairwise distances between 1,000 points — with no loop at all.
- Break it on purpose: build a
(100,)array of predictions and a(100, 1)array of targets, subtract them, and look at the shape you get. Then compute a mean-squared error from it and see how wrong the number is. This bug is worth meeting deliberately once. - From scratch: implement standardization, a stable softmax, and one-hot encoding using only array operations — no Python loops. Check your softmax against
scipy.special.softmax. - Read the shapes: open any model file in a repo you like and, for each layer, write down the input and output shapes. Fluency here is what makes reading unfamiliar ML code fast.
Try it right here — edit and run the code, and if you get stuck or hit an error, ask Ada on the right: she can see your code and your terminal output.
Ready for the next step? Take these arrays to labelled, real-world data in Pandas & DataFrames, learn to draw them in Plotting & Visualization, or see the mathematics they encode in Linear Algebra.