NumPy, pandas and the idea of tensors

Meera has been thinking. “We have six months of orders now. Three questions. Which store sells the most plum cake per day? Do rainy days really hurt sales, or does it only feel that way? And one day, could the bot predict how many croissants each store should bake tomorrow?”

The third question is a machine learning question: learning a pattern from past data to predict the future. Priya is careful not to promise it yet. “First, we need to be able to look at the data properly. Every AI model you’ll ever train or call, from a demand forecaster to a large language model, eats arrays of numbers. So let’s learn to handle those.”

Arjun downloads the export: orders_6_months.csv. Fifty-two thousand rows. He reaches for what he knows.

The export file

To follow along, install the three libraries (uv add numpy pandas matplotlib) and run this script once. It creates a realistic, deliberately messy version of the export Arjun received. You don’t need to understand every line yet; by the end of the lesson, you will.

# make_orders.py
"""Create a realistic, slightly messy export of six months of Crumb orders."""
import numpy as np
import pandas as pd

rng = np.random.default_rng(seed=42)            # fixed seed: same "random" data every run

days = pd.date_range("2026-05-01", "2026-10-31", freq="D")
rained = rng.random(len(days)) < 0.35            # about 35% of days had rain
day_weight = np.where(rained, 0.7, 1.0)          # rainy days get fewer orders

n = 52_000
day_idx = rng.choice(len(days), size=n, p=day_weight / day_weight.sum())
stores = rng.choice(["Baner", "FC Road", "Kothrud"], size=n, p=[0.35, 0.25, 0.40])
items = rng.choice(["Butter croissant", "Plum cake", "Brownie", "Kaju katli"], size=n, p=[0.5, 0.15, 0.25, 0.10])
prices = pd.Series(items).map({"Butter croissant": 85, "Plum cake": 450, "Brownie": 90, "Kaju katli": 650})
qty = rng.integers(1, 7, size=n).astype(float)
qty[rng.choice(n, size=37, replace=False)] = np.nan          # a few missing quantities

orders = pd.DataFrame({
    "order_id": [f"CR-{100000 + i}" for i in range(n)],
    "date": days[day_idx].strftime("%Y-%m-%d"),
    "store": stores,
    "item": items,
    "qty": qty,
    "price": "₹" + prices.astype(str),             # exported as text, like the real system
    "rained": rained[day_idx],
}).sort_values(["date", "order_id"])

orders.to_csv("orders_6_months.csv", index=False)
print(f"Wrote {len(orders):,} rows")

Wrote 52,000 rows

Notice we made the rain effect up: rainy days get 70% of a normal day’s orders. That’s useful for learning, because we know the right answer in advance. With real data, you never do, and we’ll come back to what that means.

Arjun’s first attempt: loops all the way down

Arjun starts with Python’s built-in csv module and the tools from the data handling lesson:

import csv

revenue = {}
with open("orders_6_months.csv", encoding="utf-8") as f:
    for row in csv.DictReader(f):                  # each row arrives as a dict of strings
        price = int(row["price"].replace("₹", ""))
        qty = int(float(row["qty"]))
        revenue[row["store"]] = revenue.get(row["store"], 0) + price * qty

print(revenue)

ValueError: could not convert string to float: ''

Thirty-seven rows have no quantity at all. In the file, the qty field is empty, so csv hands over an empty string, and float("") has nothing to convert. He adds an if to skip those rows, then gets the revenue per store. Then Meera’s first question needs revenue per store per item per day: a dictionary of dictionaries of dictionaries. The rain question needs another pass over the file. Every new question is another loop, another nested dictionary, another twenty lines.

Then Priya asks him to time something:

import time

import numpy as np

rng = np.random.default_rng(0)
prices = rng.integers(80, 700, size=1_000_000)       # one million prices
qty = rng.integers(1, 7, size=1_000_000)             # one million quantities

price_list, qty_list = prices.tolist(), qty.tolist() # the same numbers as plain lists

start = time.perf_counter()
total = 0
for p, q in zip(price_list, qty_list):
    total += p * q
loop_ms = (time.perf_counter() - start) * 1000

start = time.perf_counter()
total_np = int((prices * qty).sum())
numpy_ms = (time.perf_counter() - start) * 1000

print(total == total_np)
print(f"Python loop: {loop_ms:.0f} ms, NumPy: {numpy_ms:.1f} ms")

True
Python loop: 107 ms, NumPy: 2.8 ms

Same answer, about forty times faster on Arjun’s laptop (your numbers will differ, but the gap won’t close). A tenth of a second sounds harmless. But training a model means doing sums like this over millions of numbers, millions of times. A 40× difference is the gap between a model training overnight and training for a month.

Egg trays and baskets

Priya takes Arjun to the stores room. On one shelf, a wicker basket: eggs, a whisk, a receipt and a lemon. On the next, an egg tray: thirty identical slots, thirty eggs.

“The basket is a Python list,” she says. “It can hold anything. But to wash everything in it, you pick up each thing, check what it is, and decide how to wash it. One at a time.”

“And the tray?”

“Every slot holds the same kind of thing, the same size, in a fixed grid. So you can rinse the whole tray in one go, under one tap.”

That’s a NumPy array: a grid of values that are all the same type, packed tightly together in memory. Because Python knows every element is, say, a 64-bit whole number, it can hand the entire grid to fast, pre-compiled code that processes it in one sweep, instead of looping in Python and checking each element’s type. Doing an operation on a whole array at once is called vectorization.

Where the analogy breaks: an egg tray has a fixed number of slots, and so does an array. Adding one element means building a new array. Arrays are for crunching numbers, not for collections that grow one item at a time; lists are still right for that.

Following one small grid through NumPy

Before tackling 52,000 rows, Priya writes a tiny grid on the whiteboard: plum cakes sold at three stores over four days.

import numpy as np

sales = np.array([
    [12, 15, 9, 20],      # Baner
    [8, 7, 11, 6],        # FC Road
    [14, 18, 16, 22],     # Kothrud
])

print(sales.shape)              # (rows, columns)
print(sales.dtype)              # the one type every element shares
print(sales.sum())              # everything
print(sales.sum(axis=1))        # across each row: one total per store
print(sales.sum(axis=0))        # down each column: one total per day
print(sales.mean(axis=1))       # average per store
print(sales[2, 3])              # row 2 (Kothrud), column 3 (Sun)
print(sales[:, 3])              # every row, column 3: all stores on Sunday

(3, 4)
int64
158
[56 32 70]
[34 40 36 48]
[14.   8.  17.5]
22
[20  6 22]

Trace it:

The word that confuses everyone is axis. An axis is one dimension of the grid: axis 0 runs down the rows, axis 1 runs across the columns. sum(axis=0) means “collapse axis 0”, so the rows disappear and one value per column remains. Read it as “sum away this axis”.

Whole-grid arithmetic and broadcasting

Arithmetic works on every element at once:

import numpy as np

sales = np.array([[12, 15, 9, 20], [8, 7, 11, 6], [14, 18, 16, 22]])

print(sales * 450)                       # revenue: every cell times ₹450
print(sales > 15)                        # a grid of True/False
print((sales > 15).sum())                # True counts as 1: how many big days?
print(sales[sales > 15])                 # only the values that pass the test

day_price = np.array([450, 450, 450, 500])   # Sunday premium price
print(sales * day_price)                     # each column times its day's price

[[5400 6750 4050 9000]
 [3600 3150 4950 2700]
 [6300 8100 7200 9900]]
[[False False False  True]
 [False False False False]
 [False  True  True  True]]
4
[20 18 16 22]
[[ 5400  6750  4050 10000]
 [ 3600  3150  4950  3000]
 [ 6300  8100  7200 11000]]

sales > 15 produces a boolean mask, a grid of True and False, and sales[mask] keeps only the values where the mask is True. That’s filtering, the sieve from the data handling lesson, without a loop.

The last line is the clever one. sales has shape (3, 4) and day_price has shape (4,). NumPy broadcasts the smaller array: it treats the one row of prices as if it were copied down all three rows, then multiplies cell by cell. Broadcasting matches shapes starting from the right: the last dimension of each must be equal, or one of them must be 1.

pandas: the spreadsheet you can talk to

NumPy grids have positions, not names. Real data has column names, text, dates and missing values. pandas builds two labelled structures on top of NumPy arrays:

• A Series: one column of values with labels (like one column of a spreadsheet).
• A DataFrame: a table of named columns, each column a Series, all sharing one row index (the row labels).

Think of a DataFrame as the order spreadsheet, where each column is an egg tray. That’s why pandas is fast: column operations become NumPy operations underneath.

Loading and first look

import pandas as pd

orders = pd.read_csv("orders_6_months.csv")

print(orders.shape)
print(orders.head(3))
print(orders.dtypes)
print(orders.isna().sum()[orders.isna().sum() > 0])

(52000, 7)
    order_id        date    store              item  qty price  rained
0  CR-100238  2026-05-01    Baner         Plum cake  2.0  ₹450   False
1  CR-100462  2026-05-01  Kothrud  Butter croissant  6.0   ₹85   False
2  CR-100556  2026-05-01    Baner  Butter croissant  3.0   ₹85   False
order_id        str
date            str
store           str
item            str
qty         float64
price           str
rained         bool
dtype: object
qty    37
dtype: int64

Four commands, and you know the data’s size, what it looks like, each column’s type, and where the gaps are. Always start here. Two warning signs jump out of dtypes:

• qty is float64, not a whole number. NumPy’s whole-number arrays have no slot for “missing”, so pandas switched the whole column to floats to make room for NaN. Those are the 37 missing quantities that crashed Arjun’s loop.
• price is str, text. The export wrote "₹450", so you can’t do maths on it yet. (Older versions of pandas show object here instead of str; it means the same thing for us.)

Cleaning

Clean once, in a function, and reuse it. Save this as crumb_data.py:

# crumb_data.py
import pandas as pd


def load_orders(path: str = "orders_6_months.csv") -> pd.DataFrame:
    """Load the order export and fix its known problems."""
    orders = pd.read_csv(path)
    orders["price"] = orders["price"].str.replace("₹", "").astype(int)   # "₹450" -> 450
    orders = orders.dropna(subset=["qty"])                               # drop 37 rows with no qty
    orders["qty"] = orders["qty"].astype(int)
    orders["date"] = pd.to_datetime(orders["date"])                      # text -> real dates
    orders["revenue"] = orders["qty"] * orders["price"]                  # a new column, no loop
    return orders

Each line works on a whole column at once. .str.replace runs a string method on every value in the column. .astype(int) converts the column’s type. dropna(subset=["qty"]) removes rows where qty is missing. Whether to drop missing rows or fill them (fillna(1), say) is a judgement call: 37 of 52,000 rows is 0.07%, so dropping them barely changes any answer, and guessing a quantity could be wrong. With 20% missing, you’d need to think much harder. pd.to_datetime turns text into real dates, so you can ask for the weekday or month.

Answering Meera’s questions

Question 1: plum cake per store per day.

from crumb_data import load_orders

orders = load_orders()
days = orders["date"].nunique()                     # number of distinct days

plum = orders[orders["item"] == "Plum cake"]        # boolean mask: plum cake rows only
per_day = plum.groupby("store")["qty"].sum() / days

print(days, "days")
print(per_day.round(1).sort_values(ascending=False))

184 days
store
Kothrud    59.0
Baner      50.1
FC Road    37.2
Name: qty, dtype: float64

orders[orders["item"] == "Plum cake"] is the same boolean-mask filtering you saw in NumPy, with column names. groupby("store") splits the rows into one group per store, ["qty"].sum() adds up each group, and dividing by 184 days gives a daily rate. Kothrud sells about 59 plum cakes a day.

Question 2: does rain hurt sales?

from crumb_data import load_orders

orders = load_orders()
daily = orders.groupby(["date", "rained"]).size().reset_index(name="orders")
print(daily.head(3))
print(daily.groupby("rained")["orders"].mean().round(1))

        date  rained  orders
0 2026-05-01   False     359
1 2026-05-02   False     293
2 2026-05-03   False     305
rained
False    314.3
True     221.1
Name: orders, dtype: float64

This is a two-step aggregation (summarising many rows into fewer). First, count orders per day: groupby(["date", "rained"]).size() gives one row per day, with its order count. Then average those daily counts for rainy and dry days separately. Rainy days average 221 orders against 314 on dry days: about 30% fewer.

That matches the 0.7 we built into the fake data, which is reassuring. But with real data, Priya warns, be careful what you conclude. Pune’s rain falls mostly in June to September. If those months are also slow for other reasons (college holidays end, no festivals), “rain” might be standing in for “monsoon season”. Separating causes properly is a statistics question; for now, the honest answer to Meera is “rainy days have about 30% fewer orders, but we haven’t yet shown rain is the cause”.

A table for the store managers: revenue by store and item.

from crumb_data import load_orders

orders = load_orders()
table = orders.pivot_table(index="store", columns="item", values="revenue", aggfunc="sum")
print((table / 100_000).round(1))                  # in lakhs of rupees

item     Brownie  Butter croissant  Kaju katli  Plum cake
store                                                    
Baner       14.3              26.8        41.4       41.5
FC Road      9.9              19.8        29.6       30.8
Kothrud     16.5              31.2        46.7       48.8

A pivot table spreads one column across the top (items), another down the side (stores), and fills each cell with an aggregate (total revenue). Croissants sell the most units, but plum cake and kaju katli bring in the most money.

Loading other kinds of datasets

Data arrives in many shapes, and pandas reads most of them:

• pd.read_csv("file.csv") for comma-separated files, the most common export.
• pd.read_json("file.jsonl", lines=True) for JSON Lines: one JSON object per line, the usual format for LLM evaluation sets and fine-tuning data.
• pd.read_excel("file.xlsx") for spreadsheets (needs uv add openpyxl).
• pd.DataFrame(list_of_dicts) for data you already have in Python, like a list of API responses.

For AI work, you’ll also load public datasets from the Hugging Face Hub (Module 11) with its datasets library, which converts to pandas with one call:

from datasets import load_dataset          # uv add datasets

reviews = load_dataset("stanfordnlp/imdb", split="train")   # 25,000 movie reviews
df = reviews.to_pandas()
print(df.shape)

Library APIs change quickly. If a name has moved, check the docs linked below.

Looking at the data: a first plot

Numbers in a table hide patterns that a chart shows instantly. matplotlib is Python’s standard plotting library; pandas can draw with it directly.

import matplotlib.pyplot as plt

from crumb_data import load_orders

orders = load_orders()
croissants = orders[orders["item"] == "Butter croissant"]
daily = croissants.groupby("date")["qty"].sum()

fig, ax = plt.subplots(figsize=(10, 4))
daily.plot(ax=ax, linewidth=1)
daily.rolling(7).mean().plot(ax=ax, linewidth=2)       # 7-day average smooths the noise
ax.set_title("Butter croissants sold per day, all stores")
ax.set_xlabel("Date")
ax.set_ylabel("Croissants")
fig.tight_layout()
fig.savefig("croissants_per_day.png", dpi=120)
print("Saved croissants_per_day.png")

Saved croissants_per_day.png

Open the image. You’ll see a jagged daily line, dipping on rainy days, with a smoother line running through it. rolling(7).mean() computes a moving average: each point is the average of the last seven days, which hides day-to-day noise and shows the trend. A chart like this is the first thing to make before attempting Meera’s third question. A forecasting model is, in a sense, a precise way of drawing that smooth line one day into the future. Always label the axes and give the chart a title: a chart without them is a puzzle.

Tensors: the same idea, more dimensions

Now the idea underneath all of deep learning. Priya stacks things on the table as she talks.

“One number, the rating a customer gave: that’s a point.” She puts down one egg.

“A row of numbers, plum cakes sold at three stores today: a line.” A row of eggs.

“Our sales grid, stores by days: a flat tray.” A tray.

“A stack of trays, one grid per item, for four items: a crate.” A crate of trays.

Each step adds one dimension. A tensor is the general name for a grid of numbers with any number of dimensions. NumPy arrays are tensors; PyTorch and TensorFlow have their own tensor types that work the same way and can also run on a GPU (a graphics processor, which does thousands of simple sums at once). The number of dimensions is sometimes called the tensor’s rank.

import numpy as np

rating = np.array(4.5)
today = np.array([12, 8, 14])
grid = np.zeros((3, 4))
crate = np.zeros((4, 3, 4))
photo_batch = np.zeros((32, 224, 224, 3), dtype=np.float32)

for name, t in [("rating", rating), ("today", today), ("grid", grid), ("crate", crate), ("photos", photo_batch)]:
    print(f"{name:7} ndim={t.ndim} shape={t.shape} numbers={t.size:,}")

rating  ndim=0 shape=() numbers=1
today   ndim=1 shape=(3,) numbers=3
grid    ndim=2 shape=(3, 4) numbers=12
crate   ndim=3 shape=(4, 3, 4) numbers=48
photos  ndim=4 shape=(32, 224, 224, 3) numbers=4,816,896

A batch of 32 small photos is nearly five million numbers. That’s why vectorization isn’t optional in deep learning, and why GPUs matter.

The last row of the table is the shape you’ll meet most in this course. When an LLM reads text, each token becomes a list of numbers (an embedding, a list of numbers that captures what a piece of text means). A batch of sentences becomes a 3-D tensor: sentences × tokens × numbers-per-token. Everything a neural network does is arithmetic on tensors like this one.

Here’s a preview of Module 07 in one line of NumPy. Suppose an embedding model turned three phrases into three short vectors (real ones have hundreds of numbers). Similar meanings get similar vectors, and you measure similarity with a dot product of normalised vectors, called cosine similarity:

import numpy as np

phrases = ["eggless chocolate cake", "cake without eggs", "delivery is late"]
vectors = np.array([
    [0.9, 0.8, 0.1],
    [0.8, 0.9, 0.2],
    [0.1, 0.2, 0.9],
])

unit = vectors / np.linalg.norm(vectors, axis=1, keepdims=True)   # scale each row to length 1
similarity = unit @ unit.T                                        # every pair at once
print(similarity.round(2))

[[1.   0.99 0.3 ]
 [0.99 1.   0.39]
 [0.3  0.39 1.  ]]

@ is matrix multiplication; .T flips rows and columns. “Eggless chocolate cake” and “cake without eggs” score 0.99, nearly identical in meaning, while “delivery is late” scores only 0.3 and 0.39. That’s the core of semantic search and RAG, and it’s an operation on a 2-D tensor. You’ll build it for real in Module 07.

And when you get to PyTorch (Module 27), moving between the two worlds is one line:

import torch                                  # uv add torch

tensor = torch.from_numpy(np.zeros((3, 4)))   # NumPy array -> PyTorch tensor
array = tensor.numpy()                        # and back

Why this comes before PyTorch, TensorFlow and Hugging Face

You could jump straight to a deep-learning framework. People who do tend to get stuck in the same places:

• Shapes. Most deep-learning errors are shape errors: “expected (32, 768), got (768, 32)”. If axis, broadcasting and reshape are familiar from NumPy, those errors are readable.
• Data preparation. Before any model trains, someone loads, cleans, filters and splits the data. That’s pandas.
• Inspecting results. Model predictions and evaluation scores come back as arrays. Putting them in a DataFrame next to the inputs is how you find out where a model is wrong.
• The framework APIs copy NumPy. PyTorch’s tensor.sum(dim=1), tensor.shape and boolean masks behave like their NumPy cousins (PyTorch says dim where NumPy says axis).

Where it breaks

• Text that looks like numbers. "₹450" loads as text. A column of numbers with one stray "N/A" loads as text too, and .sum() then glues strings together or fails. Check dtypes first, every time.
• NaN behaves differently in different places. np.mean([4, np.nan, 5]) is nan: one missing value poisons the whole result. pd.Series([4, np.nan, 5]).mean() is 4.5: pandas skips missing values by default. Both are reasonable, and mixing them up gives silently different answers.
• Axis confusion. sum(axis=0) versus sum(axis=1) gives answers with the same total and different meanings. Print the shape of the result and ask, “is this one number per store, or per day?”
• Memory. pandas loads the whole file into memory, often taking several times the file’s size on disk. Fifty-two thousand rows is nothing; fifty million may not fit on a laptop. Load only the columns you need (usecols=[...]), or use a database.
• Plausible but wrong conclusions. Clean code can still answer the wrong question. Rainy days having fewer orders doesn’t prove rain causes it. Charts and group averages are where questions start, not where they end.

Arjun explains it to Meera

Arjun sends Meera a two-page summary, a pivot table and a chart, built from about twenty lines of pandas. His loop version had reached ninety lines and answered one question.

“So which store sells the most plum cake?” she asks.

“Kothrud, about 59 a day, then Baner at 50 and FC Road at 37,” Arjun says. “On rainy days we get about 30% fewer orders, 221 instead of 314. But most of the rain falls in monsoon months, so part of that drop might be the season, not the rain. I’d want to compare rainy and dry days within the same weeks before you plan baking around it.”

“And how did you get through 52,000 orders so quickly?”

“Before, I read the orders one by one, like checking every egg in a basket. Now the data sits in trays, one tray per column, and each question is one instruction to the whole tray at once: ‘only plum cake’, ‘group by store’, ‘add up’. It’s about forty times faster, and the code reads like your questions.”

“And predicting croissants?”

“That’s the next step. This chart is where it starts.”

Try it yourself

Using load_orders() from crumb_data.py:

1. Add a weekday column with orders["date"].dt.day_name().
2. Find the average number of croissants sold per day for each weekday (total croissants per date first, then average by weekday). Print it rounded to 1 decimal place, in Monday-to-Sunday order.
3. Print which weekday is the busiest for croissants.

Expected (from the seeded data):

weekday
Monday       499.2
Tuesday      500.5
Wednesday    523.0
Thursday     487.4
Friday       481.3
Saturday     498.6
Sunday       493.9
Name: qty, dtype: float64
Busiest: Wednesday

from crumb_data import load_orders

orders = load_orders()
orders["weekday"] = orders["date"].dt.day_name()

croissants = orders[orders["item"] == "Butter croissant"]
per_date = croissants.groupby(["date", "weekday"])["qty"].sum().reset_index()
by_weekday = per_date.groupby("weekday")["qty"].mean()

week = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"]
by_weekday = by_weekday.reindex(week)              # put the rows in calendar order

print(by_weekday.round(1))
print("Busiest:", by_weekday.idxmax())

weekday
Monday       499.2
Tuesday      500.5
Wednesday    523.0
Thursday     487.4
Friday       481.3
Saturday     498.6
Sunday       493.9
Name: qty, dtype: float64
Busiest: Wednesday

Common mistakes

• Looping over DataFrame rows. for i, row in orders.iterrows(): works, slowly. Almost every loop over rows can be a column operation, a mask or a groupby.
• Skipping the first look. Run shape, head(), dtypes and isna().sum() before any analysis.
• Forgetting that filters return a new DataFrame. orders[orders["store"] == "Baner"] doesn’t change orders. Assign the result to a name.
• Mixing up axis=0 and axis=1. Check the result’s shape against what you expected.
• Charts without labels. Always set a title and axis labels.

Recap

• NumPy arrays are fixed-type grids of numbers; operating on a whole array at once (vectorization) is many times faster than looping in Python.
• shape, axis, boolean masks and broadcasting are the core ideas: sum(axis=0) collapses rows, masks filter, and broadcasting stretches shapes that match from the right.
• pandas DataFrames are labelled tables built on NumPy columns: load with read_csv, inspect with head, dtypes and isna, clean column by column, and summarise with groupby and pivot_table.
• matplotlib turns summaries into charts; a moving average reveals trends behind noisy daily numbers.
• A tensor is a grid of numbers with any number of dimensions; images, batches and token embeddings are all tensors, and every neural network is arithmetic on them.

That’s the end of the Python refresher. Arjun can now read and write the Python that AI engineering runs on: data structures, classes, errors, files, types, async, HTTP, model calls and arrays. Next, Priya promises, they’ll open up the model itself. “You’ve been sending text to Claude and getting text back. Next module: what actually happens in between, and why the model sometimes confidently makes things up.”

Go deeper

• NumPy documentation, especially “NumPy: the absolute basics for beginners”.
• pandas documentation, including “10 minutes to pandas”.
• matplotlib documentation, with a gallery of example plots.
• Hugging Face documentation, for the datasets library.