Numerical & Tensor Computing

Why numpy (and not Python lists) intermediate

Here's a moment nearly every AI engineer lives through once. You prototype a similarity search with plain Python lists — a few loops, some sum() and zip() — and on a hundred vectors it's instant. Then you point it at a real corpus, a hundred thousand embeddings, and it grinds for minutes per query. Nothing about the logic changed. What changed is that you were doing million-number arithmetic in the slowest possible place: the Python interpreter. numpy exists to move that arithmetic somewhere fast, and understanding why it's faster is the difference between reaching for it reflexively and fighting it.

The plain-language reason comes down to how the numbers are stored. Picture a Python list of floats as a cloakroom full of numbered tickets: each slot in the list doesn't hold a number, it holds a ticket (a pointer) pointing off to a boxed float object living somewhere else in memory entirely. To add two lists element by element, Python must, for every single pair, follow a ticket, unbox the object, check its type, do the add, box the result, and store a new ticket. That bookkeeping — not the addition — is where all the time goes. A numpy array throws the cloakroom out: it's one solid, contiguous slab of raw numbers, all the same type, sitting shoulder-to-shoulder in memory like the array from D1.

That single design choice buys two things at once. First, memory: no per-element object headers, no scattered pointers — often 10× less RAM for the same data, which matters the moment you hold an embedding matrix for a million documents. Second, speed: because the numbers are typed and packed, numpy runs the operation as one tight compiled C loop over the buffer, with no per-element interpreter overhead, frequently using SIMD instructions that do several numbers per CPU cycle — and it releases the GIL (A2) so other threads keep working. The common mistake is to assume numpy is fast because "C is faster than Python"; the deeper truth is that the contiguous typed layout is what makes the fast C loop possible in the first place. That's why the entire ML stack — PyTorch, TensorFlow, every embedding library — is built on this one idea, and why the rest of this page is really about learning to stay in that fast layout instead of accidentally dropping back into Python loops.

1 · The ndarray & dtypes intermediate

Before you can do anything clever with numpy, you have to internalise two words, because 90% of numpy bugs are really confusion about one of them: shape and dtype. Get comfortable asking every array "what shape are you, and what type are you?" and most of the mysterious errors — the shape mismatches, the silent precision loss, the model that eats twice the RAM you budgeted — stop being mysterious.

Think of an ndarray ("n-dimensional array") as a spreadsheet that can have more than two dimensions. A 1-D array is a single row of numbers (a vector — one embedding). A 2-D array is a grid with rows and columns (a matrix — a whole batch of embeddings stacked up). Go to 3-D and you've got a stack of grids (say, a batch of images, each a grid of pixels). The shape is just the size along each of those dimensions, written as a tuple: (3,) is three numbers in a line, (1000, 384) is a thousand rows of 384 numbers each. Reading a shape out loud — "a thousand vectors, each 384-dimensional" — is a skill you'll use constantly, and it's the fastest way to catch a bug: if the shape isn't what you expected, something upstream went wrong.

The dtype is the second half of the contract: what kind of number sits in every slot — float32, float64, int8. Unlike a Python list, every element in one array shares a single dtype, and that's precisely what lets numpy pack them tightly and compute fast. The dtype also fixes the price: a float32 costs 4 bytes, a float64 costs 8, so shape × dtype-size = exactly how much memory this array eats — no guessing. That formula is worth burning in, because it's how you predict, before you allocate, whether a 7-billion-parameter model will fit on your card. The classic beginner trap here is invisible: numpy defaults to float64, but the ML world runs on float32 (and smaller), so an array you forgot to cast is quietly doubling your memory and slowing your math. One more thing to hold onto as you read the lab: slicing an array (emb[0], emb[:10]) usually gives you a view — a window onto the same memory, not a copy — which is why it's instant and free, but also why writing into a slice changes the original. Fast, and occasionally surprising.

pythonimport numpy as np

v = np.array([1.0, 2.0, 3.0], dtype=np.float32)
print(v.shape, v.dtype)          # (3,) float32

m = np.zeros((3, 4), dtype=np.float32)   # a 3x4 matrix of zeros
print(m.shape, m.nbytes)         # (3, 4)  48  (3*4*4 bytes)

# An "embedding matrix": 1000 vectors of dimension 384
emb = np.random.rand(1000, 384).astype(np.float32)
print(emb.shape, emb.nbytes // 1024, "KB")   # (1000, 384) ~1500 KB

# Reshape/slice are VIEWS (no copy) — like memoryview in A2
row = emb[0]                     # a view into row 0
block = emb[:10]                 # first 10 rows, still a view

2 · Vectorization — kill the loop advanced

If you learn one habit from this entire page, make it this one. Vectorization is the single change that most often turns "my script is too slow to use" into "my script is done before I let go of Enter" — and it's also the habit that most reliably marks the difference between someone who imported numpy and someone who actually thinks in it. The good news: it's less a technique to memorise than a reflex to build — every time your fingers start typing for i in range(len(...)) over an array, a little alarm should go off.

Here's the idea in plain terms. "Vectorizing" means describing the operation on the whole array at once — a * b + 1 — instead of walking element by element and doing the arithmetic yourself. It feels like it shouldn't matter; the same multiplications happen either way. But where they happen is everything. A Python for loop does each multiply up in the interpreter, paying that per-element unbox-check-rebox tax from the last section a million times over. Write a * b + 1 and numpy hands the entire job to a single compiled C loop that rips through the contiguous buffer, often several numbers per clock cycle via SIMD. The analogy: looping in Python is filling a swimming pool with a teaspoon; vectorizing is opening the fire hose. Same water, wildly different time.

The mental shift is to stop thinking "for each element, do X" and start reaching for the whole-array vocabulary that expresses the same intent without a loop. There are really only a handful of tools, and this lab shows the core set: elementwise math (a * b + 1 combines matching positions across arrays), reductions that crush an array down to a summary (.sum(), .mean(), .max()), and boolean masks (a[a > 0.9]), which replace filter loops by building a True/False array and keeping only the True positions. The common mistake — and it's a subtle one — is a "half-vectorized" loop that calls a numpy operation inside a Python for: you pay the interpreter tax on every iteration and get almost none of the speedup. The win only lands when the loop itself disappears into a single array expression. When you genuinely can't express something that way, that's what np.where, fancy indexing, and (last resort) np.vectorize are for — but reach for the whole-array form first, every time.

pythonimport numpy as np

a = np.random.rand(1_000_000)
b = np.random.rand(1_000_000)

# ✗ Python loop — slow (millions of interpreted iterations)
# out = [a[i]*b[i] + 1 for i in range(len(a))]

# ✓ Vectorized — one C-level call, ~100x faster
out = a * b + 1                # elementwise multiply then add
total = (a * b).sum()          # dot-product-like reduction
mean = a.mean(); mx = a.max()  # reductions are C loops too

# masking (boolean index) replaces filter loops
big = a[a > 0.9]              # all elements > 0.9, vectorized

3 · Broadcasting advanced

Broadcasting is the feature that feels like magic right up until it feels like betrayal. When it works, you add a bias vector to every row of a matrix in one clean expression and marvel at how little code it took. When it goes wrong, you get either a cryptic could not be broadcast together error or — far worse — no error at all and a silently wrong-shaped result that poisons everything downstream. Learning the rule properly is what keeps you on the magic side.

The plain-language idea: broadcasting lets numpy combine arrays of different shapes by pretending to stretch the smaller one to match the bigger one — without actually copying any data. Picture adding a single bias vector of length 384 to all 1000 rows of a matrix. You could physically tile that vector into a full 1000×384 grid and add grid-to-grid, but that wastes memory copying the same numbers a thousand times. Broadcasting is the lazy, smart version: numpy just reuses the one vector for every row, reading it 1000 times from the same place. Think of it like a stencil — one cut-out, stamped across the whole page, never redrawn.

Now the rule that decides whether two shapes are compatible, and it's worth stating precisely because "roughly right" is where the silent bugs live. numpy lines the two shapes up from the right and compares them dimension by dimension. For each pair, they're compatible only if they're equal, or one of them is 1 (the 1 is the dimension that gets stretched). So (1000, 384) with (384,) works — the vector stretches down every row; and (1000, 384) with (1000, 1) works — that column stretches across every column, which is exactly the per-row normalization in the lab. The classic mistake shows up right there: np.linalg.norm(emb, axis=1) returns shape (1000,), and lining (1000,) up from the right against (1000, 384) tries to match 1000 against 384 — mismatch, error. The fix is keepdims=True, which preserves the shape as (1000, 1) so the trailing 1 can broadcast. This is why seasoned numpy users reach for keepdims=True almost automatically, and why, when a broadcast result looks weird, the very first thing to print is the shapes of both operands — nine times out of ten the bug is a dimension that collapsed when you needed it kept.

pythonimport numpy as np

emb = np.random.rand(1000, 384).astype(np.float32)

# norms: shape (1000, 1); emb: shape (1000, 384) -> broadcasts across columns
norms = np.linalg.norm(emb, axis=1, keepdims=True)   # (1000, 1)
unit = emb / norms                                       # each row scaled to length 1
print(np.linalg.norm(unit[0]))                          # ~1.0

# Broadcasting rules: compare shapes right-to-left; dims must be equal or 1.
#   (1000, 384) with (384,)   -> ok (row added to every row)
#   (1000, 384) with (1000,1) -> ok (scale each row)

4 · Cosine similarity & top-k — the RAG core advanced

This is the section where the abstract array math suddenly becomes the thing you're actually building. Every "chat with your documents" app, every semantic search box, every RAG pipeline in this course, has this exact computation beating at its centre: given a question, find the handful of stored passages most related to it. The two previous sections — vectorization and broadcasting — were the warm-up; this is the payoff, and it's remarkably little code.

Start with the human question: how do you measure whether two pieces of text are "about the same thing" once they've been turned into vectors? The answer is cosine similarity — and the key insight is that it cares about direction, not length. Picture each embedding as an arrow pointing somewhere in space. Two arrows pointing the same way are about the same topic even if one is longer; two pointing at right angles are unrelated. Cosine similarity is just the cosine of the angle between them: 1.0 means identical direction, 0 means perpendicular/unrelated. Using the angle rather than raw distance is what makes it robust — it won't be fooled into thinking a long document and a short one are different just because their vectors have different magnitudes.

Now the trick that makes it fast enough for real corpora. Cosine is formally the dot product divided by both vectors' lengths — but if you normalize every vector to length 1 up front (the exact broadcasting move from section 3), that divisor becomes 1 and cosine similarity collapses into a plain dot product. And a dot product of your query against every one of 10,000 stored vectors is nothing more than a single matrix-vector multiply: corpus @ query, one @, all 10,000 scores at once, in C. That's the whole engine. The last piece is picking the winners: you usually want the top 5, and the beginner instinct — fully sort all 10,000 scores — is wasteful when you only need 5. np.argpartition is the right tool: it does an O(n) partition to shove the k best to the front without ordering the rest, and then you sort just those k. The common mistakes to file away: forgetting to normalize (your "cosine" scores are then meaningless), and reaching for a full argsort when argpartition would do — a detail interviewers love because it shows you know top-k isn't the same problem as sorting.

pythonimport numpy as np

def normalize(x):
    return x / np.linalg.norm(x, axis=-1, keepdims=True)

corpus = normalize(np.random.rand(10000, 384).astype(np.float32))
query  = normalize(np.random.rand(384).astype(np.float32))

# Score ALL 10k docs at once: (10000,384) @ (384,) -> (10000,)
scores = corpus @ query                    # cosine sim (already normalized)

# Top-5 without a full sort: argpartition is O(n), then sort the 5
k = 5
idx = np.argpartition(-scores, k)[:k]      # indices of the 5 best (unordered)
top = idx[np.argsort(-scores[idx])]        # order those 5
print(top, scores[top])

5 · SciPy & linear algebra intermediate

Cosine similarity handles the "dense" side of search — but half of a good retrieval system runs on a completely different kind of matrix, and if you store it the naive way you'll run out of memory before you finish loading it. That's the problem this section fixes, and it's where numpy's ecosystem partner, SciPy, earns its place. Think of SciPy as the well-stocked toolbox built on top of numpy: optimization, statistics, signal processing, and — the two things that matter most for AI engineering — sparse matrices and linear-algebra solvers.

The star here is the sparse matrix, and to see why it matters, picture the keyword side of search: a term-document matrix (the guts of TF-IDF and BM25). Rows are words in your vocabulary — maybe 50,000 of them — columns are documents, and each cell counts how often a word appears in a document. The catch: any single document uses only a few hundred distinct words, so almost every cell is zero. Storing that as a normal dense array is like mailing a 50,000-page book to describe a page-long letter — you'd pay for millions of zeros. A sparse matrix stores only the non-zero entries plus their coordinates, so its memory cost is proportional to the actual data, not the empty grid. The .nnz ("number of non-zeros") attribute in the lab is the whole point made concrete: 2 stored values where a dense array would hold 6, and at real vocabulary sizes that ratio becomes thousands-to-one. The mistake to avoid is calling .toarray() on a big sparse matrix "just to look at it" — that materializes every zero and is exactly the memory blow-up you were avoiding.

The second tool, np.linalg.solve, is the quiet workhorse under a startling amount of ML. When you write A x = b, you're saying "find the x that satisfies this whole system of linear equations at once," and solve does it directly — no rearranging by hand, no guessing. Why care? Because fitting a line (or a plane, or a regression) to data is, underneath, exactly this: least-squares regression boils down to solving a linear system. So the same one-liner that answers a toy two-equation puzzle in the lab is the machinery behind fitting models to data. The honest caveat, worth knowing before you lean on it: solve assumes the system has one clean answer (a square, non-singular A); for the over-determined "more data points than unknowns" case that real regression actually is, you reach for np.linalg.lstsq instead — same spirit, built for the messy real shape.

pythonimport numpy as np
from scipy import sparse

# A term-document matrix is mostly zeros -> store sparse (memory: O(nonzeros))
tdm = sparse.csr_matrix(np.array([[0, 2, 0], [1, 0, 0]]))
print(tdm.nnz)                   # 2 stored values, not 6

# Solve Ax = b (numpy) — the workhorse of least-squares, regressions
A = np.array([[3.0, 1], [1, 2]]); b = np.array([9.0, 8])
x = np.linalg.solve(A, b)        # [2., 3.]

6 · Tensors & automatic differentiation expert advanced

Here's the question this section quietly answers: when people say a model "learns," what is the actual mechanical thing happening? The mystique around deep learning evaporates once you see it, and what you see is smaller than you'd think — a bookkeeping trick sitting on top of the array math you already know. If you build the tiny engine in the lab, you will understand the single most important mechanism in all of modern AI, and you'll understand it because you wrote it, not because someone asserted it.

First, the word tensor, which sounds intimidating and mostly isn't. A tensor is just an n-dimensional array — numpy's ndarray — handed two extra superpowers. One: it can live on a GPU, so those big matrix multiplies run on thousands of cores at once. Two, and this is the interesting one: it remembers the operations that were applied to it. That memory is what makes automatic differentiation possible. Everything else about PyTorch and TensorFlow is engineering around those two additions.

Now the mechanism itself, in plain language. Training a model means nudging its millions of weights in whatever direction reduces the error — and to nudge each weight correctly you need its gradient: "if I wiggle this one weight a hair, how much does the final error change, and which way?" Computing that by hand for millions of weights is hopeless, so frameworks compute it automatically. The trick, called reverse-mode autograd, has two phases. On the forward pass, as you do the arithmetic, each result secretly records who its parents were and how it was made — building a graph of the whole computation. Then backward() walks that graph in reverse, applying the chain rule from calculus to hand each node its share of the gradient. The intuition for the chain rule is a bucket brigade: the error at the end is passed backward hand to hand, and at every junction it splits according to how much each input contributed. The one non-obvious detail — and it's a lovely callback — is the order you must walk the graph: every node has to be processed after all the nodes that fed into it, which is precisely a topological sort of a DAG (the algorithm from D5). The common misconception this section kills is that autograd does "symbolic calculus" or numerical finite-differences; it does neither. It just records a graph of primitive ops, each of which knows its own tiny local derivative, and multiplies those together backwards. That's it. That's how an LLM learns.

pythonclass Value:
    """A scalar that remembers how it was computed, so it can backprop."""
    def __init__(self, data, _parents=(), _backward=None):
        self.data = data
        self.grad = 0.0
        self._parents = _parents
        self._backward = _backward or (lambda: None)
    def __add__(self, other):
        out = Value(self.data + other.data, (self, other))
        def _back():
            self.grad += out.grad        # d(a+b)/da = 1
            other.grad += out.grad
        out._backward = _back
        return out
    def __mul__(self, other):
        out = Value(self.data * other.data, (self, other))
        def _back():
            self.grad += other.data * out.grad   # d(a*b)/da = b
            other.grad += self.data * out.grad
        out._backward = _back
        return out
    def backward(self):
        self.grad = 1.0
        # visit nodes in reverse topological order (D5!) and apply chain rule
        topo, seen = [], set()
        def build(v):
            if v not in seen:
                seen.add(v)
                for p in v._parents: build(p)
                topo.append(v)
        build(self)
        for v in reversed(topo):
            v._backward()

# f = a*b + a  ; df/da = b + 1, df/db = a
a, b = Value(2.0), Value(3.0)
f = a * b + a
f.backward()
print(a.grad, b.grad)            # 4.0 (=3+1)  2.0

7 · dtypes & quantization expert advanced

This is the section that decides whether that impressive open model runs on your laptop or just refuses to load. Every engineer who has tried to run a local LLM has hit the same wall: "7 billion parameters" sounds abstract until you discover it wants 28 GB of RAM you don't have. Quantization is the lever that gets it down to something a normal machine can hold — and the whole idea is built on the dtype knowledge from section 1, now used deliberately as a cost dial.

Start with the arithmetic that makes it concrete. A model's "size" is mostly its weights, and a weight's memory cost is fixed by its dtype: a float32 weight is 4 bytes. So 7 billion weights × 4 bytes ≈ 28 GB — that's the table's top row, and it's why full-precision models are a data-center thing. Halve the bits to 16 (float16 or bfloat16) and you halve the memory to ~14 GB; the useful wrinkle is that bfloat16 keeps float32's range (fewer overflow surprises) at the cost of precision, which is why training loves it. Drop to 8-bit integers and you're at ~7 GB; to 4-bit (the Q4 GGUF files people actually download) and a 7B model fits in ~3.5 GB — comfortably on a laptop. Each halving of bits roughly halves the memory and speeds up the math, because moving fewer bytes is the real bottleneck. There's no free lunch, though, and naming the tradeoff is the point: fewer bits means coarser numbers, so you're trading a little accuracy for a lot of fit-on-my-machine.

So what is quantization, mechanically? It's mapping a range of fine-grained float values onto a small set of integer buckets — the same idea as rounding temperatures to the nearest degree. The lab shows the honest version: find one scale factor so the largest-magnitude weight lands on the top integer (127 for int8), then round every weight to its nearest integer step. Storing those int8 values is 4× smaller than the floats. To use the weights you dequantize — multiply the integers back by the scale — recovering numbers that are close to, but not exactly, the originals. That gap is the quantization error, and here's the intuition that ties it all together: with a fixed number of buckets, a wider spread of values means each bucket has to cover more ground, so rounding loses more. That's why the "try this" nudges you to scale the weights up and watch the error grow, and it's why real quantization schemes work hard to keep outlier weights from stretching the range for everyone else. The naive whole-tensor scale in the lab is the teaching version; production methods (per-channel scales, GPTQ, AWQ) are refinements of exactly this move.

pythonimport numpy as np

weights = np.random.randn(1000).astype(np.float32)      # trained weights

# Symmetric int8 quantization: map [-max, max] -> [-127, 127]
scale = np.abs(weights).max() / 127
q = np.round(weights / scale).astype(np.int8)          # 4x smaller than float32
dequant = q.astype(np.float32) * scale                 # approximate original
err = np.abs(weights - dequant).mean()
print(q.nbytes, weights.nbytes, f"mean err={err:.5f}")   # 1000 vs 4000 bytes

8 · Local LLM inference expert expert

This is where every thread on the page ties together into one running program. The contiguous arrays, the vectorized matmuls, the tensors, the quantized weights — a local LLM generating text is all of those at once, in a loop. If the inner workings of ChatGPT or a local model have ever felt like a black box, this section pries the lid off, and the reveal is almost anticlimactic: it's a short loop wrapped around math you've now seen.

The mental model to hold is that an LLM does not write a sentence; it writes one token, then does the whole thing again to write the next, and again, until it decides to stop. Picture someone writing a message one keystroke at a time, re-reading everything written so far before choosing each next character — that's the loop. Concretely, four steps repeat: (1) a forward pass takes the tokens so far and, through a stack of big matrix multiplies — the exact math from sections 2–4, now running on quantized weights from section 7 — produces logits, one raw score for every possible next token; (2) softmax turns those thousands of scores into a probability distribution; (3) the model samples one token from that distribution; (4) it appends that token and loops back. Stop when it emits an end-of-sequence token or hits your length limit. Getting the model onto your machine is just the setup before this loop: download a quantized file (usually GGUF for CPU/Apple-Silicon via llama.cpp, or GPTQ/AWQ for GPU) and load it with a runtime.

The one piece of real code here is softmax, and it hides the single most-asked numerical-stability question in ML interviews, so it's worth slowing down on. Softmax exponentiates every logit and divides by the total, converting arbitrary scores into positive numbers that sum to 1. The trap: np.exp of a large logit overflows to infinity, and you get nan — garbage that silently corrupts everything after it. The fix is the line z = z - z.max(): subtract the biggest logit from all of them before exponentiating. Because that same shift appears in both the numerator and the denominator of the final ratio, it cancels perfectly — the probabilities are mathematically identical — but now the largest exponent is e^0 = 1 and nothing overflows. It's a free safety net, and knowing why it's free (the cancellation) is what interviewers are probing for. The other knob to notice is temperature, which divides the logits before softmax: high temperature flattens the distribution (more random, more "creative"), low temperature sharpens it toward the top choice (more focused, more deterministic) — the very same temperature slider the Anthropic API and every other LLM API exposes. Understanding it here means understanding what that parameter actually does everywhere else.

python# With llama-cpp-python (pip install llama-cpp-python), a quantized local model:
#
#   from llama_cpp import Llama
#   llm = Llama(model_path="model-Q4_K_M.gguf", n_ctx=4096)
#   out = llm("Explain RAG in one sentence.", max_tokens=64)
#   print(out["choices"][0]["text"])
#
# Under the hood each step is:
#   logits = forward(tokens)          # big matmuls (this page's math), on CPU/GPU
#   probs  = softmax(logits / temperature)
#   next   = sample(probs, top_p=..., top_k=...)   # pick a token
#   tokens.append(next)               # then repeat until EOS or max_tokens

import numpy as np
def softmax(logits, temp=1.0):
    z = logits / temp
    z = z - z.max()                    # numerical stability (avoid overflow)
    e = np.exp(z)
    return e / e.sum()
print(softmax(np.array([2.0, 1.0, 0.1])).round(3))   # [0.659 0.242 0.099]

Exercises expert

🎯 Interview practice interview

The interview questions this topic gets asked — worked, with code. For the full pattern catalog see A9 · Big Tech AI-engineering patterns.

pythonimport numpy as np
def softmax(logits):
    z = logits - np.max(logits)   # stability
    e = np.exp(z)
    return e / e.sum()

pythonimport numpy as np
def attention(Q, K, V):
    d = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d)
    scores -= scores.max(axis=-1, keepdims=True)
    w = np.exp(scores); w /= w.sum(axis=-1, keepdims=True)
    return w @ V

Checkpoint expert

import numpy as np

data = list(range(10))

loop = sum(x*x for x in data)          # Python loop
vec  = int((np.arange(10) ** 2).sum()) # vectorized

print(loop, vec, loop == vec)          # 285 285 True

import numpy as np

X = np.array([[3.0, 4.0], [1.0, 0.0], [0.0, 2.0]])
norms = np.linalg.norm(X, axis=1, keepdims=True)   # shape (3,1)
unit = X / norms                                    # broadcast (3,1)->(3,2)

print(np.round(unit, 3))
print(np.round(np.linalg.norm(unit, axis=1), 6))    # [1. 1. 1.]

import numpy as np

def top_k(query, docs, k=3):
    q = query / np.linalg.norm(query)
    D = docs / np.linalg.norm(docs, axis=1, keepdims=True)
    sims = D @ q                                 # (n,) cosine scores
    idx = np.argpartition(-sims, k-1)[:k]        # top-k unordered
    return idx[np.argsort(-sims[idx])], sims     # then sort those k

rng = np.random.default_rng(0)
docs = rng.random((8, 5))
q = docs[2] + 0.01                                # near doc 2
idx, sims = top_k(q, docs, k=3)
print("top-3 doc ids:", idx.tolist())            # doc 2 should lead
print("scores:", np.round(sims[idx], 3))

import numpy as np

def softmax_naive(x):
    e = np.exp(x)
    return e / e.sum()

def softmax_stable(x):
    z = x - np.max(x)          # shift-invariant, prevents overflow
    e = np.exp(z)
    return e / e.sum()

big = np.array([1000.0, 1001.0, 1002.0])
print("naive:", softmax_naive(big))    # nan/inf -- overflow
print("stable:", np.round(softmax_stable(big), 4))  # [0.09 0.2447 0.6652]
print("sums to 1:", round(float(softmax_stable(big).sum()), 6))

import numpy as np

def quantize(x):
    scale = np.abs(x).max() / 127.0        # symmetric int8
    q = np.round(x / scale).astype(np.int8)
    return q, scale

def dequantize(q, scale):
    return q.astype(np.float32) * scale

x = np.array([0.12, -0.98, 0.55, 0.03, -0.44], dtype=np.float32)
q, scale = quantize(x)
x2 = dequantize(q, scale)

print("int8:", q.tolist())
print("max abs error:", round(float(np.abs(x - x2).max()), 5))
print("bytes: %d -> %d" % (x.nbytes, q.nbytes))   # 20 -> 5

import numpy as np

class VectorIndex:
    def __init__(self):
        self.ids = []
        self.M = None
    def _norm(self, V):
        V = np.asarray(V, dtype=np.float32)
        return V / (np.linalg.norm(V, axis=1, keepdims=True) + 1e-9)
    def add(self, ids, vecs):
        V = self._norm(vecs)
        self.M = V if self.M is None else np.vstack([self.M, V])
        self.ids.extend(ids)
    def search(self, queries, k=3):
        Q = self._norm(queries)            # (q, d)
        sims = Q @ self.M.T                # (q, n) all queries at once
        k = min(k, self.M.shape[0])
        part = np.argpartition(-sims, k-1, axis=1)[:, :k]
        out = []
        for r in range(Q.shape[0]):
            cols = part[r][np.argsort(-sims[r, part[r]])]
            out.append([(self.ids[c], round(float(sims[r, c]), 3)) for c in cols])
        return out

rng = np.random.default_rng(1)
idx = VectorIndex()
idx.add([f"d{i}" for i in range(6)], rng.random((6, 4)))
print(idx.search(rng.random((2, 4)), k=2))   # top-2 (id, score) per query

Knowledge check check yourself