ML Coding RoundsIntermediate20 min read5 questions
NumPy, pandas & ML Algorithms From Scratch
The ML coding round is not LeetCode. It is 'implement k-means in 20 minutes' and 'vectorise this loop' — with an interviewer watching how you debug.
Covers: Vectorisation, broadcasting, pandas performance, implementing k-means/kNN/linear regression, gradient descent by hand
ML coding rounds test three things: can you write clean vectorised NumPy, can you implement a standard algorithm from its definition, and can you debug your own code without panicking. They are far more predictable than general algorithm interviews — the set of things they ask is small and learnable.
Filter
0/5 mastered
IntermediatenumpyvectorizationbroadcastingAsked at Meta, Amazon
30-second answer
Use the expansion ‖a−b‖² = ‖a‖² + ‖b‖² − 2a·b so the whole distance matrix is one matmul plus two broadcasts. Broadcasting aligns shapes from the right, and dimensions are compatible if they are equal or one of them is 1; size-1 dimensions are stretched without copying. This turns an O(n·m) Python loop into a single BLAS call — typically 100–1000× faster.
pythonThree versions, with timings
import numpy as np, time
A = np.random.default_rng(0).normal(size=(2000, 128))
B = np.random.default_rng(1).normal(size=(1500, 128))
# 1. Nested loop -- what candidates write under pressure
def dist_loop(A, B):
out = np.empty((len(A), len(B)))
for i in range(len(A)):
for j in range(len(B)):
out[i, j] = np.sqrt(((A[i] - B[j]) ** 2).sum())
return out
# 2. Broadcasting -- correct, but allocates an (n, m, d) intermediate
def dist_broadcast(A, B):
diff = A[:, None, :] - B[None, :, :] # (2000, 1500, 128) = 3 GB!
return np.sqrt((diff ** 2).sum(-1))
# 3. Matmul identity -- the answer they want
def dist_matmul(A, B):
a2 = (A ** 2).sum(1)[:, None] # (n, 1)
b2 = (B ** 2).sum(1)[None, :] # (1, m)
d2 = a2 + b2 - 2 * (A @ B.T) # broadcasts to (n, m)
return np.sqrt(np.maximum(d2, 0)) # clamp: fp error can give -1e-14
for name, fn in [("matmul", dist_matmul), ("broadcast", dist_broadcast)]:
t = time.perf_counter(); fn(A, B); print(f"{name:10s} {time.perf_counter()-t:.4f}s")
t = time.perf_counter(); dist_loop(A[:200], B[:200])
print(f"{'loop':10s} {(time.perf_counter()-t)*(2000*1500)/(200*200):.4f}s (extrapolated)")
# matmul 0.0121s
# broadcast 1.4830s <- 120x slower AND uses 3 GB of RAM
# loop 9.8412s <- 800x slower
#
# Note the np.maximum(d2, 0): the identity can produce tiny negatives from
# floating-point cancellation, and sqrt of a negative gives NaN. Interviewers
# specifically look for this guard.Broadcasting rules
Align shapes from the right
(5, 3) and (3,) align as (5, 3) and (1, 3). Missing leading dimensions are treated as 1.Dimensions are compatible if equal or one is 1
(5, 3) and (1, 3) → OK. (5, 3) and (5, 1) → OK. (5, 3) and (4, 3) → error.Size-1 dimensions are stretched without copying
NumPy uses stride-0 views, so broadcasting itself is free. The result still allocates, which is why version 2 above uses 3 GB.
pythonCommon broadcasting patterns worth having memorised
import numpy as np
X = np.random.default_rng(0).normal(size=(100, 5))
# Standardise columns
Xs = (X - X.mean(0)) / X.std(0) # (100,5) - (5,) -> (5,) broadcasts
# Row-wise normalisation -- note keepdims, the classic bug
Xn = X / np.linalg.norm(X, axis=1, keepdims=True) # (100,5) / (100,1)
# Without keepdims you get (100,5) / (100,) -> ShapeError or silent wrong answer
# One-hot encoding without a loop
y = np.array([0, 2, 1, 2, 0])
onehot = (y[:, None] == np.arange(3)[None, :]).astype(float)
# Pairwise cosine similarity
def cosine_matrix(A, B):
An = A / np.linalg.norm(A, axis=1, keepdims=True)
Bn = B / np.linalg.norm(B, axis=1, keepdims=True)
return An @ Bn.T
# Numerically stable softmax over rows
def softmax(Z):
Z = Z - Z.max(axis=1, keepdims=True) # subtract max BEFORE exp
E = np.exp(Z)
return E / E.sum(axis=1, keepdims=True)
# Top-k indices per row without sorting everything
def topk(scores, k):
idx = np.argpartition(-scores, k, axis=1)[:, :k] # O(n) not O(n log n)
rows = np.arange(len(scores))[:, None]
order = np.argsort(-scores[rows, idx], axis=1)
return idx[rows, order]Say these points
- ‖a−b‖² = ‖a‖² + ‖b‖² − 2a·b turns pairwise distance into one matmul.
- Broadcasting aligns right; dims must be equal or 1; size-1 dims stretch via stride-0.
- Guard sqrt with np.maximum(d2, 0) — fp cancellation produces tiny negatives.
- The 3D difference tensor is a memory trap, not a speed trap.
- keepdims=True is required whenever you divide by a reduction.
Avoid these mistakes
- Forgetting keepdims and getting a silent shape error.
- sqrt of a slightly negative squared distance producing NaN.
- Building an (n, m, d) intermediate on large inputs.
Expect these follow-ups
- How would you compute pairwise distances for 10M points?
- Why is np.argpartition faster than np.argsort for top-k?
- How does numpy decide whether an operation copies or returns a view?
IntermediatekmeansclusteringimplementationAsked at Amazon, Apple
30-second answer
Initialise centroids with k-means++ (pick the first uniformly, then each subsequent one with probability proportional to squared distance from the nearest existing centroid), then alternate assignment (each point to its nearest centroid) and update (each centroid to the mean of its points) until the assignments stop changing or the inertia change falls below a tolerance.
pythonComplete implementation with the edge cases handled
import numpy as np
def _pairwise_sq_dist(X, C):
"""(n, k) squared distances via the matmul identity."""
d2 = (X ** 2).sum(1)[:, None] + (C ** 2).sum(1)[None, :] - 2 * (X @ C.T)
return np.maximum(d2, 0)
def kmeans_plusplus(X, k, rng):
"""Spread initial centroids: D^2 sampling. Reduces bad local minima a lot."""
n = len(X)
centroids = [X[rng.integers(n)]]
closest_sq = _pairwise_sq_dist(X, np.array(centroids)).ravel()
for _ in range(1, k):
probs = closest_sq / closest_sq.sum()
centroids.append(X[rng.choice(n, p=probs)])
new_sq = _pairwise_sq_dist(X, centroids[-1][None, :]).ravel()
closest_sq = np.minimum(closest_sq, new_sq) # keep distance to NEAREST
return np.array(centroids)
def kmeans(X, k, max_iter=300, tol=1e-4, n_init=10, seed=0):
rng = np.random.default_rng(seed)
best = (None, None, np.inf)
for _ in range(n_init): # multiple restarts: k-means is
C = kmeans_plusplus(X, k, rng) # non-convex, restarts matter
prev_inertia = np.inf
for _ in range(max_iter):
d2 = _pairwise_sq_dist(X, C)
labels = d2.argmin(1)
inertia = float(d2[np.arange(len(X)), labels].sum())
new_C = np.empty_like(C)
for j in range(k):
mask = labels == j
if mask.any():
new_C[j] = X[mask].mean(0)
else:
# EMPTY CLUSTER: reseed at the point furthest from its centroid.
# Forgetting this gives NaN centroids -- the classic bug.
new_C[j] = X[d2[np.arange(len(X)), labels].argmax()]
C = new_C
if abs(prev_inertia - inertia) <= tol * abs(prev_inertia):
break
prev_inertia = inertia
if inertia < best[2]:
best = (labels.copy(), C.copy(), inertia)
return {"labels": best[0], "centroids": best[1], "inertia": best[2]}
# Verify against sklearn
from sklearn.cluster import KMeans
X = np.random.default_rng(0).normal(size=(600, 4)) + \
np.repeat(np.array([[0,0,0,0],[6,6,0,0],[0,6,6,0]]), 200, axis=0)
mine = kmeans(X, 3, seed=0)
ref = KMeans(3, n_init=10, random_state=0).fit(X)
print(f"mine inertia={mine['inertia']:.3f} sklearn inertia={ref.inertia_:.3f}")
# mine inertia=2381.442 sklearn inertia=2381.442| Property | Value |
|---|---|
| Time per iteration | O(n·k·d) |
| Space | O(n·k) for the distance matrix, O(k·d) for centroids |
| Convergence | Guaranteed to a local minimum (inertia decreases monotonically) |
| Global optimum | NP-hard; hence n_init restarts |
| Assumptions | Spherical, similar-sized clusters; Euclidean distance |
Say these points
- Alternate assignment and update until assignments stabilise or inertia plateaus.
- k-means++ samples each new centroid with probability ∝ D² to the nearest existing one.
- Handle empty clusters explicitly or you get NaN centroids.
- Use n_init restarts — the objective is non-convex.
- O(n·k·d) per iteration; use mini-batch k-means at large n.
Avoid these mistakes
- No empty-cluster handling.
- Random initialisation with a single run.
- Nested Python loops for the distance computation.
Expect these follow-ups
- How would you pick k without labels?
- How does mini-batch k-means change the update rule?
- Why does k-means struggle with elongated or nested clusters?
Intermediatelogistic-regressiongradient-descentimplementationAsked at Google, Two Sigma
30-second answer
Use a numerically stable sigmoid and a log-loss computed via logaddexp so you never take log(0). The gradient is Xᵀ(p − y)/n plus the regularisation term, with the intercept excluded from the penalty. Add a convergence check on the gradient norm and support mini-batching.
pythonNumerically stable implementation
import numpy as np
def stable_sigmoid(z):
"""Avoids overflow: exp(710) is inf in float64."""
out = np.empty_like(z, dtype=float)
pos, neg = z >= 0, z < 0
out[pos] = 1.0 / (1.0 + np.exp(-z[pos]))
ez = np.exp(z[neg])
out[neg] = ez / (1.0 + ez)
return out
def log_loss_from_logits(z, y):
"""log(1+exp(z)) - y*z, computed via logaddexp -> never log(0)."""
return float(np.mean(np.logaddexp(0, z) - y * z))
class LogisticRegressionGD:
def __init__(self, lr=0.1, n_iter=1000, l2=0.0, l1=0.0,
batch_size=None, tol=1e-6, seed=0):
self.lr, self.n_iter = lr, n_iter
self.l2, self.l1 = l2, l1
self.batch_size, self.tol = batch_size, tol
self.rng = np.random.default_rng(seed)
def fit(self, X, y):
n, d = X.shape
self.w = np.zeros(d)
self.b = 0.0
self.history = []
for it in range(self.n_iter):
if self.batch_size:
idx = self.rng.choice(n, min(self.batch_size, n), replace=False)
Xb, yb = X[idx], y[idx]
else:
Xb, yb = X, y
z = Xb @ self.w + self.b
p = stable_sigmoid(z)
err = p - yb # the clean gradient term
grad_w = Xb.T @ err / len(yb)
grad_b = err.mean() # intercept: NOT penalised
if self.l2:
grad_w += self.l2 * self.w
if self.l1:
grad_w += self.l1 * np.sign(self.w)
self.w -= self.lr * grad_w
self.b -= self.lr * grad_b
loss = log_loss_from_logits(X @ self.w + self.b, y)
if self.l2: loss += 0.5 * self.l2 * (self.w ** 2).sum()
if self.l1: loss += self.l1 * np.abs(self.w).sum()
self.history.append(loss)
gnorm = np.sqrt((grad_w ** 2).sum() + grad_b ** 2)
if gnorm < self.tol:
break
return self
def predict_proba(self, X):
return stable_sigmoid(X @ self.w + self.b)
def predict(self, X, threshold=0.5):
return (self.predict_proba(X) >= threshold).astype(int)
# Validate against sklearn on the same objective
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import make_classification
X, y = make_classification(n_samples=2000, n_features=8, random_state=0)
mine = LogisticRegressionGD(lr=0.5, n_iter=5000, l2=1/2000).fit(X, y)
ref = LogisticRegression(C=1.0, max_iter=5000).fit(X, y)
print("mine :", np.round(mine.w, 3))
print("sklearn:", np.round(ref.coef_[0], 3))
print("accuracy:", (mine.predict(X) == y).mean(), (ref.predict(X) == y).mean())
# mine : [ 0.096 -0.03 0.048 1.905 0.012 -0.014 -0.024 0.031]
# sklearn: [ 0.097 -0.03 0.048 1.907 0.012 -0.014 -0.024 0.031]
# accuracy: 0.9105 0.9105pythonGradient check — always do this for a hand-rolled model
def gradient_check(model, X, y, eps=1e-6, n_checks=5):
"""Finite differences vs analytic gradient. Catches sign and factor errors."""
rng = np.random.default_rng(0)
z = X @ model.w + model.b
p = stable_sigmoid(z)
analytic = X.T @ (p - y) / len(y) + model.l2 * model.w
for _ in range(n_checks):
j = rng.integers(len(model.w))
orig = model.w[j]
model.w[j] = orig + eps
loss_up = log_loss_from_logits(X @ model.w + model.b, y) \
+ 0.5 * model.l2 * (model.w ** 2).sum()
model.w[j] = orig - eps
loss_dn = log_loss_from_logits(X @ model.w + model.b, y) \
+ 0.5 * model.l2 * (model.w ** 2).sum()
model.w[j] = orig
numeric = (loss_up - loss_dn) / (2 * eps)
rel = abs(numeric - analytic[j]) / max(abs(numeric), abs(analytic[j]), 1e-12)
print(f"w[{j}] analytic={analytic[j]:+.8f} numeric={numeric:+.8f} rel_err={rel:.2e}")
assert rel < 1e-5
# w[3] analytic=+0.00012417 numeric=+0.00012417 rel_err=3.11e-09Say these points
- Branch the sigmoid on the sign of z to avoid exp overflow.
- Compute log-loss from logits with logaddexp — never log(p).
- Gradient is Xᵀ(p − y)/n; the intercept is never penalised.
- Always gradient-check a hand-rolled model with finite differences.
- Extensions: softmax multi-class, class weights, IRLS, early stopping.
Avoid these mistakes
- Naive sigmoid producing inf/nan on large |z|.
- Regularising the intercept.
- No convergence criterion, so it runs the full n_iter regardless.
Expect these follow-ups
- Extend this to multi-class softmax regression.
- How would Newton's method change the convergence rate and the cost per step?
- How do you handle a feature that perfectly separates the classes?
Advancedpandasperformancedata-engineeringAsked at Instacart, Airbnb
30-second answer
Profile first, then apply in order: eliminate `.apply` and `iterrows` in favour of vectorised operations or `groupby.transform`, downcast dtypes and use categoricals (often 5–10× memory reduction), avoid repeated concatenation in loops, use merge on indexed and sorted keys, chunk or move to a columnar engine (Polars, DuckDB) when the data no longer fits comfortably in memory.
pythonThe four transformations that give most of the speedup
import pandas as pd, numpy as np
# ---------- 1. Kill .apply and iterrows ----------
# SLOW: a Python function call per row
df["margin"] = df.apply(lambda r: (r.revenue - r.cost) / r.revenue
if r.revenue > 0 else 0, axis=1) # ~180 s @ 50M
# FAST: vectorised with np.where
df["margin"] = np.where(df.revenue > 0, (df.revenue - df.cost) / df.revenue, 0)
# ~0.6 s (300x)
# ---------- 2. groupby.transform instead of merge-back ----------
# SLOW: groupby, then merge the aggregate back
agg = df.groupby("user_id")["amount"].mean().reset_index()
df = df.merge(agg, on="user_id", suffixes=("", "_mean")) # ~45 s
# FAST: transform broadcasts the aggregate in place
df["amount_mean"] = df.groupby("user_id")["amount"].transform("mean") # ~8 s
# ---------- 3. Dtypes: the biggest memory win ----------
def optimise_dtypes(df):
for col in df.select_dtypes("object"):
if df[col].nunique() / len(df) < 0.5:
df[col] = df[col].astype("category") # often 10x smaller
for col in df.select_dtypes("integer"):
df[col] = pd.to_numeric(df[col], downcast="integer")
for col in df.select_dtypes("float"):
df[col] = pd.to_numeric(df[col], downcast="float")
return df
print(f"before: {df.memory_usage(deep=True).sum()/1e9:.2f} GB")
df = optimise_dtypes(df)
print(f"after : {df.memory_usage(deep=True).sum()/1e9:.2f} GB")
# before: 14.82 GB
# after : 2.91 GB <- 5x, and every subsequent operation gets faster too
# ---------- 4. Never concat in a loop ----------
# SLOW: O(n^2) -- each concat copies the whole frame
out = pd.DataFrame()
for chunk in chunks:
out = pd.concat([out, process(chunk)]) # quadratic
# FAST: collect, concat once
out = pd.concat([process(c) for c in chunks], ignore_index=True)| Operation | Slow form | Fast form | Typical speedup |
|---|---|---|---|
| Row-wise logic | df.apply(..., axis=1) | np.where / vectorised ops | 100–500× |
| Iteration | for _, row in df.iterrows() | vectorised or itertuples | 50–100× |
| Group aggregate back to rows | groupby + merge | groupby.transform | 5–10× |
| String ops | .apply(lambda s: ...) | .str accessor | 10–50× |
| Repeated lookup | df[df.id == x] in a loop | set index, use .loc | 100×+ |
| Building a frame | concat in a loop | list then one concat | O(n²) → O(n) |
| Category columns | object dtype | category dtype | 5–10× memory |
pythonThe escape hatches
# --- DuckDB: SQL over Parquet, no full load, out-of-core ---
import duckdb
result = duckdb.sql("""
SELECT user_id,
COUNT(*) AS n_txn,
SUM(amount) AS total,
AVG(amount) AS avg_amount,
MAX(ts) - MIN(ts) AS span
FROM 'events/*.parquet'
WHERE ts >= '2026-01-01'
GROUP BY user_id
HAVING COUNT(*) > 5
""").df() # only the RESULT enters pandas
# --- Polars: lazy, multi-threaded, query-optimised ---
import polars as pl
out = (pl.scan_parquet("events/*.parquet") # lazy: nothing read yet
.filter(pl.col("ts") >= pl.datetime(2026, 1, 1))
.group_by("user_id")
.agg([pl.len().alias("n_txn"),
pl.col("amount").sum().alias("total"),
pl.col("amount").mean().alias("avg_amount")])
.filter(pl.col("n_txn") > 5)
.collect()) # optimised plan runs here
# --- Chunked processing when you must stay in pandas ---
agg = {}
for chunk in pd.read_csv("huge.csv", chunksize=1_000_000):
for uid, grp in chunk.groupby("user_id"):
s = agg.setdefault(uid, {"n": 0, "total": 0.0})
s["n"] += len(grp); s["total"] += grp.amount.sum()Say these points
- `.apply(axis=1)` and `iterrows` are the usual culprits — vectorise them.
- `groupby.transform` beats groupby-then-merge for broadcasting aggregates.
- Categorical and downcast dtypes often cut memory 5×, speeding everything up.
- Never concat inside a loop — it is quadratic.
- Beyond ~100M rows, move to Polars or DuckDB over Parquet.
Avoid these mistakes
- Optimising before profiling.
- Leaving high-cardinality strings as object dtype.
- Merging on unsorted, unindexed keys repeatedly.
Expect these follow-ups
- How does Polars' lazy execution enable optimisations pandas cannot do?
- How would you compute a rolling 30-day per-user aggregate over 500M rows?
- When is a merge better done in SQL than in pandas?
AdvancedpytorchattentionimplementationAsked at OpenAI, Anthropic
30-second answer
Project to Q, K, V, reshape to (batch, heads, seq, head_dim), compute scaled dot-product scores with a causal mask, softmax, multiply by V, then merge heads and apply the output projection. For generation, cache K and V across steps and only compute Q, K, V for the new token — turning each step from O(n²) to O(n).
pythonMulti-head attention with KV caching
import math, torch, torch.nn as nn, torch.nn.functional as F
class MultiHeadAttention(nn.Module):
def __init__(self, d_model=512, n_heads=8, dropout=0.0):
super().__init__()
assert d_model % n_heads == 0
self.n_heads = n_heads
self.head_dim = d_model // n_heads
self.qkv = nn.Linear(d_model, 3 * d_model, bias=False) # fused projection
self.out = nn.Linear(d_model, d_model, bias=False)
self.dropout = dropout
def forward(self, x, kv_cache=None, causal=True):
B, T, C = x.shape
# (B, T, 3C) -> 3 x (B, n_heads, T, head_dim)
q, k, v = self.qkv(x).split(C, dim=2)
q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
if kv_cache is not None:
if kv_cache.get("k") is not None:
k = torch.cat([kv_cache["k"], k], dim=2)
v = torch.cat([kv_cache["v"], v], dim=2)
kv_cache["k"], kv_cache["v"] = k, v
# With a cache and T == 1, every cached position is in the past,
# so no mask is needed for the single new query row.
causal = causal and T > 1
scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
if causal:
Tq, Tk = q.size(-2), k.size(-2)
mask = torch.ones(Tq, Tk, dtype=torch.bool, device=x.device).tril(
diagonal=Tk - Tq)
scores = scores.masked_fill(~mask, float("-inf"))
attn = F.softmax(scores, dim=-1)
if self.dropout and self.training:
attn = F.dropout(attn, p=self.dropout)
y = attn @ v # (B, H, Tq, head_dim)
y = y.transpose(1, 2).contiguous().view(B, T, C) # merge heads
return self.out(y)
# --- verify against PyTorch's fused kernel ---
torch.manual_seed(0)
mha = MultiHeadAttention(64, 4).eval()
x = torch.randn(2, 10, 64)
with torch.no_grad():
mine = mha(x)
B, T, C = x.shape
q, k, v = mha.qkv(x).split(C, dim=2)
q, k, v = (t.view(B, T, 4, 16).transpose(1, 2) for t in (q, k, v))
ref = mha.out(F.scaled_dot_product_attention(q, k, v, is_causal=True)
.transpose(1, 2).reshape(B, T, C))
print("max diff:", float((mine - ref).abs().max())) # max diff: 2.4e-07pythonGeneration with and without the cache — the speedup
import time, torch
@torch.no_grad()
def generate(model, prompt_ids, max_new=200, use_cache=True):
ids = prompt_ids
caches = [{} for _ in model.blocks] if use_cache else None
for step in range(max_new):
if use_cache:
# Only the NEW token goes through the model after the first step
inp = ids if step == 0 else ids[:, -1:]
else:
inp = ids # recompute everything
logits = model(inp, caches)
next_id = logits[:, -1].argmax(-1, keepdim=True)
ids = torch.cat([ids, next_id], dim=1)
return ids
for cache in (False, True):
t0 = time.perf_counter()
generate(model, prompt, max_new=200, use_cache=cache)
print(f"use_cache={cache}: {time.perf_counter()-t0:.2f}s")
# use_cache=False: 8.41s <- O(n^2) total work
# use_cache=True : 0.62s <- O(n) total work, ~13x faster at 200 tokens
#
# The gap widens with sequence length: at 2000 tokens it is ~100x.Say these points
- Fuse the QKV projection; reshape to (B, H, T, head_dim); scale by √head_dim.
- Mask before softmax with −inf, and offset the mask correctly when using a cache.
- KV caching turns generation from O(n²) to O(n) — ~13× at 200 tokens.
- `.contiguous()` is required after transpose before `.view()`.
- Use `F.scaled_dot_product_attention` in production for the fused Flash kernel.
Avoid these mistakes
- Incorrect causal mask when the query length is 1 and the cache is long.
- Recomputing K and V for the whole prefix at every generation step.
- Forgetting to divide by √head_dim.
Expect these follow-ups
- Extend this to grouped-query attention.
- How would you add RoPE to this implementation?
- How does the memory footprint of the KV cache grow, and how do you bound it?
Showing 5 of 5 questions for numpy-pandas-ml-coding-interview.
Check your understanding
5 questions · no sign-up, nothing stored
Hands-on challenge
Build it — this is what you talk about in a deep-dive round.
Build a mini ML library and benchmark it
Implement the core algorithms from scratch, validated against reference implementations.
Requirements
- k-means with k-means++ and empty-cluster handling; match sklearn's inertia exactly on the same seed.
- Logistic regression with stable sigmoid and logaddexp loss, L1/L2, and a passing gradient check.
- kNN with the matmul distance identity, supporting 100k×100k without exceeding memory.
- Multi-head attention with a KV cache, verified against `F.scaled_dot_product_attention` to 1e-5.
Stretch goals
- Add mini-batch k-means and compare convergence and wall-clock against the full version at n = 5M.
- Profile a real pandas pipeline, then rewrite the hot path in Polars or DuckDB and report the speedup.