Knowledge BaseFoundations

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.

beginner#python#numpy#tooling

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: .

The whole page in two sentences

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.

How to read this page

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.

Why this is a prerequisite for almost everything

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.

Think of it like a shopping bag versus an egg carton:

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] * 2 does not double your numbers — it gives you a longer list, [1, 2, 3, 1, 2, 3].
Doubling three numbers, both ways

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:

  1. Start an empty result list.
  2. Take 1, multiply by 2, get 2, append it.
  3. Take 2, multiply by 2, get 4, append it.
  4. Take 3, multiply by 2, get 6, append it.
  5. 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]) * 2array([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:

Python · runs in your browser
What this does: Shows the single most important behavioural difference between a Python list and a NumPy array: multiplying a list repeats it, while multiplying an array does real arithmetic on every element at once. This one line is the difference between writing loops forever and writing vectorized code.
Try to recall

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.

  1. What shape is it? How big is the grid, along each direction?
  2. What dtype is it? What kind of number is in each slot?
  3. 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 .

Reading four shapes you will actually meet
  1. (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.
  2. (3, 4) — a table with 3 rows and 4 columns. Two dimensions, 12 numbers total. Row count first, always.
  3. (256, 256, 3) — a colour image: 256 pixels tall, 256 wide, and 3 numbers per pixel (red, green, blue). 196,608 numbers.
  4. (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.

An image is just an array — an 8 by 8 grid of numbers— interactive, drag & zoom
Loading chart…
Each square is one element of an array of shape (8, 8). The background ramps from dark on the left to lighter on the right because its values run 20, 28, 36 and so on; the bright block is a 3 by 3 patch whose values were set to 240. Nothing here is an image file — it is arithmetic you can slice, add and multiply. Hover any cell to read its stored number.

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:

dtypeBytes eachWhat it's for
float648NumPy's default for decimals. Maximum precision, maximum memory.
float324The deep-learning standard. Half the memory, plenty of precision for gradients.
float16 / bfloat162Mixed-precision training on GPUs — fast, and delicate.
int648Default for whole numbers on Linux and macOS (Windows historically defaulted to 32-bit).
bool1Masks — the True/False arrays you filter with.
Why anyone cares about 4 bytes versus 8

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.

Memory to hold one billion parameters, by dtype— interactive, drag & zoom
Loading chart…
The same billion numbers, stored three ways. float64 is NumPy's default and costs 8 GB; float32 - what PyTorch defaults to - costs 4 GB; float16 costs 2 GB and is why mixed-precision training exists. This is a straight multiplication, not a benchmark: bytes per element times one billion.

Now inspect some real arrays — shape, dtype, dimensions and the memory they occupy:

Python · runs in your browser
What this does: Creates several arrays and prints the three facts you should always know about each one - its shape, its dtype and its total memory - so you can connect the abstract idea of a grid to the concrete block of bytes it occupies. Notice that asking for float32 instead of float64 halves nbytes exactly.

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.

Stop telling the computer how, start telling it what

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.

Think of it like ordering lunch for a hundred people:

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.

Turning a written-out loop into one line

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:

  1. Make an empty results list.
  2. For each index i: take pred[i], subtract actual[i], square it, append it.
  3. Repeat for all N examples.

Vectorized version, thinking in whole arrays:

  1. diff = pred - actual — one subtraction, done to every element at once.
  2. sq = diff ** 2 — one squaring, done to every element at once.
  3. 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.

The cost of looping in Python — measured, log scale— interactive, drag & zoom
Loading chart…
Time to compute a dot product of two arrays of n float64 numbers: a hand-written Python for-loop versus NumPy. Best of five runs on this lesson's build machine, so treat the exact milliseconds as indicative and the gap as the point - at a million elements the loop takes about 37.8 ms and NumPy about 0.19 ms, roughly 200 times faster. Both axes are logarithmic; a straight line means time grows in proportion to n.

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:

Python · runs in your browser
What this does: Times the same dot product computed two ways - a hand-written Python loop versus NumPy - and prints the speedup. Run it, then change n to 500000 and run again: the gap widens, because the loop pays Python interpreter overhead once per element while NumPy pays it once per array.
Try to recall

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.

Stretch the small one, without actually copying it

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.

Think of it like a rubber stamp:

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.

Broadcasting, worked one axis at a time

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 0axis 1
col31
row (padded)14
result34

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:

0123
00123
1010111213
2020212223

Twelve additions, one line of code, zero loops, and no copy of either input was ever made.

That table, drawn:

Broadcasting a column against a row— interactive, drag & zoom
Loading chart…
A column of shape (3, 1) holding 0, 10, 20 added to a row of shape (4,) holding 0, 1, 2, 3. Neither input was copied or looped over - NumPy stretched each size-1 axis by re-reading the same memory, and produced this (3, 4) grid in one operation. Hover a cell to see the sum that landed there.

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:

Python · runs in your browser
What this does: Centres a small dataset by subtracting each column's mean, using broadcasting instead of a loop, then shows the two ways forgetting keepdims goes wrong - the loud way, where NumPy refuses with a ValueError, and the silent way, where a (4,) minus a (4,1) quietly produces a 4x4 grid of every pairwise difference instead of the 4 numbers you wanted.

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.
Cropping a patch out of an image, by hand

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.

  1. Rows first, columns second — the convention matches the shape tuple: patch = img[0:64, 0:64].
  2. Check the shape before trusting it: patch.shape is (64, 64). If it came out (64,) you indexed with one index instead of two.
  3. Average it with a reduction over both axes at once: patch.mean() — a single number.
  4. 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.

A slice is a window, not a photocopy

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().

Python · runs in your browser
What this does: Demonstrates the single most surprising NumPy behaviour - a slice shares memory with the original array, so writing to the slice edits the original - and then shows how .copy() and fancy indexing give you an independent array instead. The .base attribute tells you which kind you are holding.
Try to recall

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.

The axis you name is the axis that disappears

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.

Both reductions on the same little table, by hand

Take a (2, 3) array — 2 examples, 3 features:

f0f1f2
example 0123
example 1456
  1. A.sum(axis=0) — kill axis 0 (the examples). Add down each column: 1+4=51+4 = 5, 2+5=72+5 = 7, 3+6=93+6 = 9. Result [5, 7, 9], shape (3,) — one total per feature.
  2. A.sum(axis=1) — kill axis 1 (the features). Add across each row: 1+2+3=61+2+3 = 6, 4+5+6=154+5+6 = 15. Result [6, 15], shape (2,) — one total per example.
  3. A.sum() — no axis named, so kill them all: 2121, a single number of shape ().

Sanity check that costs nothing: both of the first two results add up to 21.

Try to recall

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?

Python · runs in your browser
What this does: Standardizes a small dataset so every feature has mean 0 and standard deviation 1, using axis=0 reductions and broadcasting - the exact preprocessing step in front of most classical ML models. The printed before-and-after summaries show three features on wildly different scales being brought onto a common one.

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.

Python · runs in your browser
What this does: Runs a complete forward pass of a one-layer classifier over a batch of 5 examples - matrix multiply, bias broadcast, numerically stable softmax, and accuracy - without a single Python loop. The final check proves that subtracting the max before exponentiating leaves the probabilities unchanged while preventing overflow.

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 and axis (renamed dim) all behave the same way.
  • A pandas DataFrame is a set of NumPy arrays with labels bolted on; df.values hands 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 NN examples beats NN 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 it yourself

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.

Recap — the key ideas
  • 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 axis argument names the dimension that disappears; keepdims=True leaves a 1 behind so the result broadcasts back.
  • A neural network layer is X @ W + b followed by a stable softmax — every idea above, in four lines and zero loops.

Practice — and how to make it stick

Learn it the way that actually works

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.

  1. 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 with np.broadcast_shapes. Every one you get wrong is a bug you would have shipped.
  2. 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.
  3. 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.
  4. 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.
  5. 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.

Practice lab
Your task: Vectorize the distance calculation. The loop version below is correct but slow - replace the TODO with a single broadcasting expression that computes all 1000 distances at once, then check that both answers agree and see how much faster it got. Bonus: once that works, compute the full 1000x1000 pairwise distance matrix using points[:, None, :] and points[None, :, :].
editor
terminal
Press Run (⌘/Ctrl+Enter) to execute.
Ask Ada — she can read your terminal

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.

Key papers