Intermediate24 min

NumPy arrays and vectors

Arrays, shapes, broadcasting, and cosine similarity — the numeric layer under every ML and embedding workflow.

Why this matters in AI / ML / GenAI

Embeddings are vectors, batches are matrices, and similarity search is a dot product. NumPy semantics — shape, dtype, broadcasting, vectorization — carry directly into PyTorch tensors. Shape mismatch is the single most common ML error message.

Why arrays beat lists

A NumPy array is a fixed-type block of memory. Operations run in compiled code over the whole array at once — often 10 to 100 times faster than a Python loop.

Two attributes you check constantly:

  • shape — a tuple, e.g. (32, 768) = 32 rows of a 768-dimensional embedding
  • dtype — float32 for models, float64 by default in NumPy, int64 for token ids

Mixing dtypes silently upcasts and doubles memory. Model code standardises on float32.

Vectorization and broadcasting

Vectorization means expressing computation over whole arrays: a * 2 + b rather than looping element by element.

Broadcasting lets arrays of different shapes combine when dimensions are compatible (equal, or one of them is 1). Subtracting a (768,) mean vector from a (32, 768) batch works: the mean is applied to every row.

Rules, right to left: dimensions must be equal or one must be 1. (32, 768) with (768,) is fine. (32, 768) with (32,) fails — you must reshape to (32, 1).

Cosine similarity

Vector search ranks by cosine similarity: the dot product of two normalized vectors.

Normalize each vector to unit length (divide by its L2 norm), then a dot product gives a value in [-1, 1] where 1 means identical direction.

Every vector database — FAISS, Pinecone, pgvector, Chroma — implements this at scale. Writing it once in NumPy demystifies the whole retrieval step.

Copy-paste examples

Copy into your own editor, or load one into the compiler below and press Run.

Shapes and dtypes

First run downloads NumPy into the browser sandbox — give it a few seconds.

import numpy as np

batch = np.random.rand(4, 8).astype(np.float32)
print("shape:", batch.shape)
print("dtype:", batch.dtype)
print("ndim:", batch.ndim, "| total values:", batch.size)
print("first row:", np.round(batch[0], 3))
print("column means:", np.round(batch.mean(axis=0), 3))
print("row means:", np.round(batch.mean(axis=1), 3))

Vectorization vs a Python loop

Same result, very different speed at scale.

import numpy as np
import time

values = np.random.rand(200_000).astype(np.float32)

start = time.perf_counter()
loop_total = 0.0
for v in values.tolist():
    loop_total += v * 2
loop_ms = (time.perf_counter() - start) * 1000

start = time.perf_counter()
vector_total = float((values * 2).sum())
vector_ms = (time.perf_counter() - start) * 1000

print(f"loop      : {loop_total:.2f} in {loop_ms:.1f} ms")
print(f"vectorized: {vector_total:.2f} in {vector_ms:.1f} ms")

Broadcasting a mean vector

Centering a batch of embeddings — one line, no loop.

import numpy as np

embeddings = np.array([
    [1.0, 2.0, 3.0],
    [2.0, 4.0, 6.0],
    [3.0, 6.0, 9.0],
], dtype=np.float32)

mean_vec = embeddings.mean(axis=0)
centered = embeddings - mean_vec

print("mean vector:", mean_vec)
print("centered:\n", centered)
print("shapes:", embeddings.shape, mean_vec.shape, centered.shape)

Cosine similarity search

This is exactly what a vector database does, minus the indexing.

import numpy as np

def normalize(matrix):
    norms = np.linalg.norm(matrix, axis=1, keepdims=True)
    return matrix / np.clip(norms, 1e-9, None)

docs = np.array([
    [0.9, 0.1, 0.0],   # python
    [0.1, 0.9, 0.0],   # kubernetes
    [0.8, 0.2, 0.1],   # python tooling
], dtype=np.float32)
labels = ["python", "kubernetes", "python tooling"]

query = np.array([[0.85, 0.15, 0.0]], dtype=np.float32)

scores = (normalize(query) @ normalize(docs).T)[0]
order = np.argsort(-scores)

for rank, idx in enumerate(order, start=1):
    print(f"{rank}. {labels[idx]:16s} score={scores[idx]:.4f}")

Rank documents against a query vector

Try it — in-browser Python

Packages: numpy

Change the query values and watch the ranking reorder.

Output

Python runs in your browser. First run downloads the runtime.

Press Run (or Ctrl+Enter) to execute.

CPython in WebAssembly. Stdlib works. NumPy and pandas load on demand. No input(), no GPU, no network installs.

Takeaways

  • Check .shape and .dtype first when debugging numeric code.
  • Vectorize instead of looping; broadcasting applies a smaller array across a larger one.
  • Cosine similarity on normalized vectors is the core of semantic search.