Part II Bank 4 NumPy

ML Coding from Scratch

Write the classic models in plain NumPy. No sklearn, no PyTorch. Show that you know the math well enough to type it.

Many Applied Science loops have one round where you build a model by hand. The interviewer wants proof that you understand what the library does for you. This page covers the ten problems that come up most. Each one has a full solution that runs, a line-by-line read, and the follow-ups you will face.

Every solution here was run and tested. Copy one into a file and run it. The test at the bottom of each file should print ok.

How ML coding rounds work
The NumPy you must know cold
Linear regression by gradient descent
Logistic regression
K-means
K-nearest neighbours
Stable softmax and cross-entropy
Best decision-tree split
ROC-AUC from scores
Self-attention forward pass
Two-layer MLP with backprop
K-fold cross-validation

How ML coding rounds work

The round is 45 to 60 minutes. You get one problem, sometimes two. You code in a shared editor or a notebook. Most of the time you may use NumPy but nothing that fits the model for you.

Minutes 0 to 5. Ask about input shapes, dtypes and edge cases. Say the math out loud before you type.
Minutes 5 to 30. Write a clean, working version. Name your shapes in comments.
Minutes 30 to 40. Test it on a tiny case you can check by hand. Fix bugs out loud.
Minutes 40 to 55. Follow-ups. Scale, numerical safety, a variant, or a gradient question.

They grade four things. Is the math right? Is the code vectorized? Is it numerically safe? Can you test and explain it? A loop over rows is a yellow flag. An exp that overflows is a red one.

Shapes first, code second. Write the shape of every array in a comment as you go. Most bugs in this round are shape bugs. A comment like # (n, k) catches them before you run anything.

The NumPy you must know cold

You will not have time to look these up. Drill them until they are automatic.

Broadcasting. NumPy lines shapes up from the right. A dim of size 1 stretches to match. So (n, 1) + (1, k) gives (n, k). Use [:, None] to add a size-1 axis.
axis. The axis you name is the one that disappears. X.sum(0) on (n, d) sums over rows and gives (d,). X.sum(1) gives (n,).
keepdims. keepdims=True keeps the reduced axis as size 1. Then the result still broadcasts back. Softmax needs this: z - z.max(1, keepdims=True).
argmax and argmin. They return the index, not the value. Ties go to the first index. Use argpartition for the top k when you do not need them sorted. It runs in O(n), not O(n log n).
Fancy indexing. P[np.arange(n), y] picks one entry per row. This is how you grab the true-class probability for cross-entropy.

einsum basics. Name each axis with a letter. Letters missing from the output get summed. einsum("ij,jk->ik") is matmul. einsum("ij,ij->i") is a row-wise dot product. einsum("bhtd,bhsd->bhts") is attention scores.
@ and batched matmul. A @ B multiplies the last two axes and broadcasts the rest. So (B, H, T, d) @ (B, H, d, T) works as is.
Stable math helpers. np.logaddexp(0, z) gives log(1 + e^z) with no overflow. np.log1p and np.expm1 stay exact near zero.
Random numbers. Use rng = np.random.default_rng(seed). Then call rng.normal, rng.permutation and rng.choice. Seeded code makes your tests repeat.
Solve, do not invert. np.linalg.solve(A, b) is faster and more stable than np.linalg.inv(A) @ b.

The silent bug. Shapes (n,) and (n, 1) broadcast into (n, n). No error is raised. Your loss just becomes wrong. Keep targets as (n,) and check shapes with an assert when in doubt.

The ten problems

1. Linear Regression by Gradient Descent Easy

The prompt

“You get a feature matrix X of shape (n, d) and targets y of shape (n,). Fit a linear model with gradient descent on mean squared error. Return the weights and the bias. Then show me the closed form.”

What they are testing

Can you derive the MSE gradient without notes?
Do you write it as one matrix product, not a loop over rows?
Do you know the normal equations and when not to use them?
Do you know why feature scaling matters for gradient descent?

Approach

The model is y_hat = X w + b. The loss is L = (1/2n) Σ (y_hat - y)². The half cancels the 2 from the square.
The gradient is dL/dw = Xᵀ (y_hat - y) / n and dL/db = mean(y_hat - y).
Step against the gradient. Stop after a fixed count or when the loss stops changing.
The closed form sets the gradient to zero. That gives the normal equations XᵀX θ = Xᵀy. Add a column of ones to learn the bias.

Solution

import numpy as np


def fit_linear_gd(X: np.ndarray, y: np.ndarray, lr: float = 0.1,
                  epochs: int = 1000, tol: float = 1e-8) -> tuple[np.ndarray, float]:
    """Fit y ~ X @ w + b by batch gradient descent on mean squared error.

    X: (n, d) features. y: (n,) targets. Returns (w of shape (d,), b).
    """
    n, d = X.shape
    w = np.zeros(d)
    b = 0.0
    prev_loss = np.inf
    for _ in range(epochs):
        err = X @ w + b - y                  # (n,) residuals
        loss = (err @ err) / (2 * n)         # half MSE, so gradients are clean
        grad_w = X.T @ err / n               # (d,)
        grad_b = err.mean()
        w -= lr * grad_w
        b -= lr * grad_b
        if abs(prev_loss - loss) < tol:      # stop when loss stops moving
            break
        prev_loss = loss
    return w, b


def fit_linear_closed(X: np.ndarray, y: np.ndarray, l2: float = 0.0) -> tuple[np.ndarray, float]:
    """Normal equations with an intercept column. Uses solve, never inv."""
    n, d = X.shape
    Xb = np.hstack([X, np.ones((n, 1))])     # (n, d+1)
    reg = l2 * np.eye(d + 1)
    reg[-1, -1] = 0.0                        # do not shrink the intercept
    theta = np.linalg.solve(Xb.T @ Xb + reg, Xb.T @ y)
    return theta[:-1], float(theta[-1])


if __name__ == "__main__":
    rng = np.random.default_rng(0)
    X = rng.normal(size=(500, 3))
    y = X @ np.array([2.0, -1.0, 0.5]) + 3.0 + 0.01 * rng.normal(size=500)
    mu, sd = X.mean(0), X.std(0)             # scale features before GD
    w, b = fit_linear_gd((X - mu) / sd, y)
    w_raw, b_raw = w / sd, b - (mu / sd) @ w  # map back to raw units
    w2, b2 = fit_linear_closed(X, y)
    assert np.allclose(w_raw, w2, atol=1e-3) and abs(b_raw - b2) < 1e-3
    print("ok", np.round(w2, 3), round(b2, 3))

Line by line

err = X @ w + b - y. One matmul gives all n residuals. The scalar b broadcasts across them.
X.T @ err / n. This is the full gradient for every weight at once. Shape (d, n) @ (n,) gives (d,).
The tol check. It stops early once the loss moves less than the tolerance. That saves time on easy problems.
np.linalg.solve. It solves the system directly. Calling inv first is slower and loses precision.
reg[-1, -1] = 0. Ridge should shrink the weights, not the intercept. So the bias entry gets no penalty.
The test scales features. GD runs on scaled data. Then the weights map back to raw units and match the closed form.

GD timeO(epochs · n d)Closed formO(n d² + d³)SpaceO(n d)

Follow-ups they will ask

When do you pick GD over the closed form? When d is large. Solving costs O(d³). GD costs O(nd) per step and also streams in mini-batches.
Why scale features? Uneven scales make the loss surface a long, thin valley. GD then zigzags and needs a tiny learning rate.
What if XᵀX is singular? Features are collinear. Add ridge (l2 > 0) or use np.linalg.lstsq, which uses the SVD.
How do you pick the learning rate? Try a log grid like 1, 0.1, 0.01. If the loss blows up, cut it. The safe bound is 2 over the largest eigenvalue of XᵀX/n.
Make it stochastic. Shuffle each epoch. Loop over mini-batches and use the same gradient on each batch.

Common traps

Forgetting to divide by n. The step size then grows with the data size.
Using y of shape (n, 1) with X @ w of shape (n,). The residual becomes (n, n).
Calling np.linalg.inv. Interviewers notice.
Penalizing the bias in ridge regression.

Say this out loud: “The gradient of half MSE is X transpose times the residual over n. I will vectorize that and check the loss goes down. For the closed form I will solve the normal equations with solve, not inv. I will mention ridge in case the matrix is singular.”

2. Logistic Regression Easy

The prompt

“Implement binary logistic regression. Input X is (n, d) and y is (n,) of zeros and ones. Train it with gradient descent and an L2 penalty. Give me fit, predict_proba and predict. Make sure it does not overflow.”

What they are testing

Do you know the log-loss gradient is Xᵀ(p - y)/n?
Can you write a sigmoid that is safe for large inputs?
Do you add L2 to the weights but not the bias?
Can you write the loss without taking log(0)?

Approach

The model is p = σ(X w + b). The loss is the mean of -y log p - (1 - y) log(1 - p).
The gradient of that loss with respect to the logit z is simply p - y. So the weight gradient matches linear regression with p in place of y_hat.
L2 adds λ w to the weight gradient.
A naive 1/(1+exp(-z)) overflows for very negative z. Split by sign so exp only sees non-positive input.

Solution

import numpy as np


def sigmoid(z: np.ndarray) -> np.ndarray:
    """Stable sigmoid. Never calls exp on a large positive number."""
    out = np.empty_like(z, dtype=float)
    pos = z >= 0
    out[pos] = 1.0 / (1.0 + np.exp(-z[pos]))
    ez = np.exp(z[~pos])                     # z < 0, so ez is in (0, 1)
    out[~pos] = ez / (1.0 + ez)
    return out


class LogisticRegression:
    def __init__(self, lr: float = 0.1, epochs: int = 2000, l2: float = 0.0) -> None:
        self.lr, self.epochs, self.l2 = lr, epochs, l2
        self.w: np.ndarray | None = None
        self.b: float = 0.0

    def fit(self, X: np.ndarray, y: np.ndarray) -> "LogisticRegression":
        """X: (n, d). y: (n,) of 0/1 labels."""
        n, d = X.shape
        self.w = np.zeros(d)
        self.b = 0.0
        for _ in range(self.epochs):
            p = sigmoid(X @ self.w + self.b)     # (n,)
            err = p - y                          # dL/dz for log loss
            grad_w = X.T @ err / n + self.l2 * self.w
            grad_b = err.mean()                  # no penalty on the bias
            self.w -= self.lr * grad_w
            self.b -= self.lr * grad_b
        return self

    def loss(self, X: np.ndarray, y: np.ndarray) -> float:
        z = X @ self.w + self.b
        # log(1 + e^z) - y*z, computed with logaddexp so it never overflows
        nll = np.mean(np.logaddexp(0.0, z) - y * z)
        return float(nll + 0.5 * self.l2 * self.w @ self.w)

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        p1 = sigmoid(X @ self.w + self.b)
        return np.column_stack([1.0 - p1, p1])   # (n, 2), like sklearn

    def predict(self, X: np.ndarray, threshold: float = 0.5) -> np.ndarray:
        return (self.predict_proba(X)[:, 1] >= threshold).astype(int)


if __name__ == "__main__":
    rng = np.random.default_rng(1)
    X = rng.normal(size=(400, 2))
    y = (X @ np.array([3.0, -2.0]) + 0.5 + 0.3 * rng.normal(size=400) > 0).astype(int)
    m = LogisticRegression(lr=0.5, l2=1e-3).fit(X, y)
    acc = (m.predict(X) == y).mean()
    assert acc > 0.95, acc
    assert np.allclose(sigmoid(np.array([-1000.0, 0.0, 1000.0])), [0.0, 0.5, 1.0])
    assert np.allclose(m.predict_proba(X).sum(1), 1.0)
    print("ok acc", acc, "loss", round(m.loss(X, y), 4))

Line by line

The two-branch sigmoid. For z ≥ 0 it uses 1/(1+e^-z). For z < 0 it uses e^z/(1+e^z). Both forms are equal, and each exp stays at or below 1.
err = p - y. This is the whole chain rule for log loss through the sigmoid. Say this fact out loud. It shows you know the math.
self.l2 * self.w. This is the gradient of (λ/2)||w||². The bias stays out.
np.logaddexp(0.0, z). It computes log(1 + e^z) safely. Log loss equals log(1 + e^z) - y z, so no log(0) ever happens.
predict_proba. It returns both columns, like sklearn. Rows sum to one.

TimeO(epochs · n d)PredictO(n d)SpaceO(n + d)

Follow-ups they will ask

Why not use MSE? With a sigmoid, MSE is not convex in w. Its gradient also vanishes when the model is very wrong. Log loss is convex.
What happens on separable data with no L2? The weights grow forever. The loss keeps dropping toward zero. L2 keeps them finite.
Extend to K classes. Use a (d, K) weight matrix and softmax. The gradient becomes Xᵀ(P - Y_onehot)/n.
How do you handle class imbalance? Weight each row's loss by its class weight. Or tune the threshold on a validation set.
Faster solver? Newton's method or L-BFGS. Newton needs the Hessian Xᵀ S X where S = diag(p(1-p)).

Common traps

Writing np.log(p) directly. It returns -inf when p rounds to 0.
Clipping p to [1e-15, 1 - 1e-15] and calling that stable. It works, but logaddexp is the cleaner answer.
Putting L2 on the bias.
Forgetting the sign. You must subtract the gradient.

Say this out loud: “The gradient through sigmoid and log loss collapses to p minus y. I will split the sigmoid by sign so exp never overflows. For the loss I will use logaddexp, so there is no log of zero.”

3. K-means Medium

The prompt

“Implement k-means. Input X is (n, d) and an integer k. Use k-means++ to start. Return the centers (k, d), the labels (n,) and the inertia. Handle empty clusters.”

What they are testing

Can you compute all point-to-center distances with no Python loop over points?
Do you know why k-means++ beats random starts?
Do you have a clear stop rule?
Do you notice that a cluster can end up with no points?

Approach

Init with k-means++. Pick the first center at random. Pick each next center with chance proportional to its squared distance from the nearest center so far.
Assign. Build the (n, k) distance matrix. Take argmin over axis 1.
Update. Each center moves to the mean of its points.
Stop when no center moves more than tol, or after max_iter rounds.
Empty cluster. Move its center to the point that is farthest from its own center.

Solution

import numpy as np


def sq_dists(X: np.ndarray, C: np.ndarray) -> np.ndarray:
    """(n, d) and (k, d) -> (n, k) squared Euclidean distances."""
    d2 = (X * X).sum(1)[:, None] + (C * C).sum(1)[None, :] - 2.0 * X @ C.T
    return np.maximum(d2, 0.0)               # clip tiny negatives from round-off


def kmeans_pp_init(X: np.ndarray, k: int, rng: np.random.Generator) -> np.ndarray:
    n = X.shape[0]
    centers = [X[rng.integers(n)]]
    d2 = sq_dists(X, centers[0][None, :])[:, 0]
    for _ in range(1, k):
        total = d2.sum()
        if total == 0:                       # all points sit on centers already
            idx = rng.integers(n)
        else:
            idx = rng.choice(n, p=d2 / total)  # pick far points more often
        centers.append(X[idx])
        d2 = np.minimum(d2, sq_dists(X, X[idx][None, :])[:, 0])
    return np.array(centers)


def kmeans(X: np.ndarray, k: int, max_iter: int = 300, tol: float = 1e-6,
           seed: int = 0) -> tuple[np.ndarray, np.ndarray, float]:
    """Lloyd's algorithm. Returns (centers (k, d), labels (n,), inertia)."""
    rng = np.random.default_rng(seed)
    C = kmeans_pp_init(X, k, rng)
    for _ in range(max_iter):
        D = sq_dists(X, C)                   # (n, k)
        labels = D.argmin(1)
        new_C = np.empty_like(C)
        for j in range(k):
            members = X[labels == j]
            if len(members) == 0:
                # Empty cluster: move it to the point worst served right now.
                far = D[np.arange(len(X)), labels].argmax()
                new_C[j] = X[far]
                labels[far] = j
                D[far, :] = 0.0              # so a second empty cluster picks another point
            else:
                new_C[j] = members.mean(0)
        shift = np.sqrt(((new_C - C) ** 2).sum(1)).max()
        C = new_C
        if shift < tol:
            break
    D = sq_dists(X, C)
    labels = D.argmin(1)
    inertia = float(D[np.arange(len(X)), labels].sum())
    return C, labels, inertia


if __name__ == "__main__":
    rng = np.random.default_rng(2)
    true = np.array([[0, 0], [5, 5], [0, 5]], dtype=float)
    X = np.vstack([c + 0.3 * rng.normal(size=(100, 2)) for c in true])
    C, labels, inertia = kmeans(X, 3)
    # Each true center should have a found center within 0.15.
    assert np.sqrt(sq_dists(true, C)).min(1).max() < 0.15
    assert np.bincount(labels).tolist() == [100, 100, 100]
    # Duplicate points force an empty cluster. It must still return k centers.
    C2, l2, _ = kmeans(np.array([[0.0, 0.0]] * 5 + [[1.0, 1.0]]), 3)
    assert C2.shape == (3, 2) and np.isfinite(C2).all()
    print("ok inertia", round(inertia, 2))

Line by line

sq_dists. It expands ||x - c||² into three terms. The cross term is one matmul. np.maximum(..., 0) removes tiny negative values from round-off.
rng.choice(n, p=d2 / total). This is the k-means++ step. Far points are more likely to become centers.
d2 = np.minimum(...). It keeps the distance to the nearest chosen center. Each new center costs O(nd), not O(nkd).
The empty-cluster branch. It takes the worst-served point and gives it to the empty cluster. Zeroing that row stops a second empty cluster from taking the same point.
shift. The largest distance any center moved. Under tol means we have converged.
Final pass. It reassigns labels to the final centers, so labels and inertia match the centers you return.

TimeO(iters · n k d)InitO(n k d)SpaceO(n k)

Follow-ups they will ask

Does k-means always converge? Yes. Each step never raises the inertia, and there are only finitely many labelings. It may stop at a local minimum.
How do you pick k? Elbow plot of inertia, silhouette score, or a downstream metric. Say the business need often sets k.
Why k-means++? It spreads the starting centers out. It gives an O(log k) bound on the expected cost. In practice it also converges faster.
Scale to a billion points? Mini-batch k-means. Update centers from small random batches with a decaying step size.
When does k-means fail? Non-round clusters, very uneven sizes, or features on different scales. Scale features first. Try a GMM or DBSCAN for odd shapes.

Common traps

Building X[:, None, :] - C[None]. That makes an (n, k, d) array and can run out of memory.
Taking the mean of an empty slice. NumPy returns nan and every later step breaks.
Stopping when labels stop changing but never capping the loop.
Forgetting to seed the random generator. Tests then fail at random.

Say this out loud: “I will use the norm expansion so the distance matrix is one matmul. I will start with k-means++ and stop when centers stop moving. If a cluster goes empty I will reseed it at the worst-fit point.”

4. K-Nearest Neighbours Medium

The prompt

“Write a k-NN classifier. Training data is X_train (n, d) with integer labels y_train (n,). Predict labels for X_test (m, d). No loops over test points. Tell me how you break ties.”

What they are testing

Do you know the distance-matrix trick ||a||² + ||b||² - 2a·b?
Do you use argpartition instead of a full sort?
Do you have a real rule for vote ties?
Can you talk about the cost of k-NN at serving time?

Approach

Build the full (m, n) squared-distance matrix with one matmul.
Use argpartition to get the k smallest per row in O(n).
Count votes per class with a one-hot sum.
Break vote ties by the smaller total distance. This is a common and easy rule to defend.

Solution

import numpy as np


def pairwise_sq_dists(A: np.ndarray, B: np.ndarray) -> np.ndarray:
    """||a - b||^2 = ||a||^2 + ||b||^2 - 2 a.b for all pairs. (m, d), (n, d) -> (m, n)."""
    a2 = np.einsum("ij,ij->i", A, A)[:, None]   # (m, 1)
    b2 = np.einsum("ij,ij->i", B, B)[None, :]   # (1, n)
    return np.maximum(a2 + b2 - 2.0 * A @ B.T, 0.0)


def knn_predict(X_train: np.ndarray, y_train: np.ndarray, X_test: np.ndarray,
                k: int = 5) -> np.ndarray:
    """Classify each test row by majority vote of its k nearest train rows.

    Ties in the vote go to the class with the smallest total distance.
    y_train holds integer labels 0..C-1.
    """
    D = pairwise_sq_dists(X_test, X_train)                  # (m, n)
    k = min(k, X_train.shape[0])
    idx = np.argpartition(D, k - 1, axis=1)[:, :k]          # k nearest, unordered
    nd = np.sqrt(np.take_along_axis(D, idx, axis=1))        # (m, k) distances
    nl = y_train[idx]                                       # (m, k) labels
    n_classes = int(y_train.max()) + 1
    onehot = np.eye(n_classes)[nl]                          # (m, k, C)
    votes = onehot.sum(1)                                   # (m, C)
    dist_sum = (onehot * nd[:, :, None]).sum(1)             # (m, C)
    # Rank by votes first, then by smaller distance. Scale keeps votes dominant.
    score = votes - dist_sum / (dist_sum.max() + 1.0) / (k + 1)
    score[votes == 0] = -np.inf
    return score.argmax(1)


if __name__ == "__main__":
    A = np.array([[0.0, 0.0], [1.0, 0.0]])
    B = np.array([[0.0, 1.0], [3.0, 4.0], [1.0, 0.0]])
    brute = ((A[:, None, :] - B[None, :, :]) ** 2).sum(-1)
    assert np.allclose(pairwise_sq_dists(A, B), brute)
    Xtr = np.array([[0.0], [1.0], [10.0], [11.0]])
    ytr = np.array([0, 0, 1, 1])
    assert knn_predict(Xtr, ytr, np.array([[0.4], [10.6]]), k=3).tolist() == [0, 1]
    # Tie: k=2, one vote each. Class 1 point is closer, so class 1 wins.
    assert knn_predict(np.array([[0.0], [3.0]]), np.array([1, 0]), np.array([[1.0]]), k=2).tolist() == [1]
    print("ok")

Line by line

einsum("ij,ij->i", A, A). Row-wise squared norms. It avoids making a temporary A * A array.
a2 + b2 - 2 A @ B.T. Broadcasting (m, 1) + (1, n) gives (m, n). Memory is O(mn), not O(mnd).
np.argpartition(D, k - 1, axis=1)[:, :k]. It puts the k smallest in the first k slots. They are not sorted, and we do not need them sorted.
take_along_axis. It gathers the matching distances per row.
np.eye(n_classes)[nl]. One-hot of shape (m, k, C). Summing over axis 1 gives votes per class.
The score line. Votes are whole numbers. The distance term is always under 1, so it only matters when votes tie.

TimeO(m n d)SelectO(m n)SpaceO(m n)TrainO(1)

Follow-ups they will ask

The distance matrix does not fit in memory. Process test rows in chunks of a few thousand. The result is the same.
Serving is too slow. Use an index. KD-trees work in low dimensions. For high dimensions use approximate search like HNSW or IVF with product quantization.
Why scale features? Distance adds up raw units. A feature in dollars swamps one in years.
Turn it into regression. Average the neighbour targets. Or weight them by 1/distance.
How do you pick k? Cross-validation. Small k overfits. Large k blurs the boundary. Odd k avoids ties in binary tasks.

Common traps

Skipping the clip at zero. Round-off can make a distance slightly negative, and sqrt then gives nan.
Using argsort on every row. It works but costs O(n log n).
Using np.bincount in a Python loop per row. That breaks the “no loops” rule.
Letting argmax break ties by class index without saying so.

Say this out loud: “I will expand the squared distance so the whole matrix is one matmul plus broadcasting. Then argpartition gets the k nearest in linear time. Ties in the vote go to the class with the smaller total distance.”

5. Stable Softmax and Cross-Entropy Easy

The prompt

“Given logits of shape (n, C) and integer labels (n,), return the mean cross-entropy loss and its gradient with respect to the logits. It must work for logits like 1000.”

What they are testing

Do you subtract the max before exp?
Do you compute log-softmax directly, not log(softmax)?
Do you know the gradient is softmax - onehot?
Do you remember to divide the gradient by n?

Approach

Softmax does not change if you add a constant to every logit. So subtract the row max. Now the biggest value is exp(0) = 1.
Log-sum-exp: log Σ e^z = m + log Σ e^(z - m). Then log-softmax is z - lse.
The loss is minus the log-prob of the true class, averaged.
The gradient for one row is p - onehot(y). For a mean over the batch, divide by n.

Solution

import numpy as np


def log_softmax(z: np.ndarray, axis: int = -1) -> np.ndarray:
    """log softmax with log-sum-exp. Subtracting the max makes exp safe."""
    m = z.max(axis=axis, keepdims=True)
    lse = m + np.log(np.exp(z - m).sum(axis=axis, keepdims=True))
    return z - lse


def softmax(z: np.ndarray, axis: int = -1) -> np.ndarray:
    e = np.exp(z - z.max(axis=axis, keepdims=True))
    return e / e.sum(axis=axis, keepdims=True)


def cross_entropy(logits: np.ndarray, y: np.ndarray) -> tuple[float, np.ndarray]:
    """Mean CE over a batch and its gradient w.r.t. the logits.

    logits: (n, C). y: (n,) integer class ids. Returns (loss, dlogits (n, C)).
    """
    n = logits.shape[0]
    logp = log_softmax(logits, axis=1)
    loss = -logp[np.arange(n), y].mean()
    grad = np.exp(logp)                      # softmax probabilities
    grad[np.arange(n), y] -= 1.0             # softmax - onehot
    return float(loss), grad / n             # divide by n because loss is a mean


if __name__ == "__main__":
    z = np.array([[1000.0, 1001.0, 1002.0], [-1000.0, 0.0, 1000.0]])
    p = softmax(z)
    assert np.all(np.isfinite(p)) and np.allclose(p.sum(1), 1.0)
    rng = np.random.default_rng(3)
    L = rng.normal(size=(4, 5))
    y = np.array([0, 3, 1, 4])
    loss, g = cross_entropy(L, y)
    eps, num = 1e-6, np.zeros_like(L)
    for i in range(L.shape[0]):
        for j in range(L.shape[1]):
            Lp, Lm = L.copy(), L.copy()
            Lp[i, j] += eps
            Lm[i, j] -= eps
            num[i, j] = (cross_entropy(Lp, y)[0] - cross_entropy(Lm, y)[0]) / (2 * eps)
    assert np.allclose(g, num, atol=1e-7)
    print("ok loss", round(loss, 4))

Line by line

m = z.max(axis, keepdims=True). keepdims keeps the shape (n, 1), so z - m broadcasts per row.
z - lse. The log-probs. This form never takes the log of a tiny number.
logp[np.arange(n), y]. Fancy indexing picks the true-class log-prob from each row.
grad[np.arange(n), y] -= 1. This turns softmax into softmax minus one-hot in place.
The test. It compares the gradient to centered finite differences. Do this in the interview if you have time.

TimeO(n C)SpaceO(n C)

Follow-ups they will ask

Derive the gradient. The answer is dL/dz_j = p_j - [j = y]. The loss is -z_y + lse(z). The first term gives -1 on the true class. The second term gives p_j.
Add label smoothing. Replace the one-hot with (1 - ε) on the true class and ε/C everywhere. The gradient becomes p - target.
Add temperature. Divide logits by T first. The gradient gets an extra factor of 1/T.
Why do frameworks fuse softmax and CE? The fused form is stable and its gradient is cheap. Separate ops would need the full Jacobian.

Common traps

np.log(softmax(z)). A probability can round to 0 and give -inf.
Subtracting the global max instead of the row max. One huge row then breaks the others.
Forgetting keepdims. (n, C) - (n,) raises an error or broadcasts wrong when n == C.
Returning the gradient without dividing by n.

Say this out loud: “Softmax is shift-invariant, so I subtract the row max first. I compute log-softmax with log-sum-exp, never log of softmax. The gradient is softmax minus one-hot, divided by the batch size.”

6. Best Decision-Tree Split Medium

The prompt

“Given numeric features X (n, d) and class labels y (n,), find the single best split. Return the feature, the threshold and the impurity gain. Support Gini and entropy. Do better than trying every threshold from scratch.”

What they are testing

Do you know the Gini and entropy formulas?
Do you know weighted child impurity and gain?
Can you get O(n log n) per feature with sort and cumulative counts?
Do you handle tied feature values and empty children?

Approach

Gini is 1 - Σ p_c². Entropy is -Σ p_c log₂ p_c.
Gain is parent impurity minus the size-weighted average of the two children.
The naive way tries n thresholds and rescans n rows each time. That is O(n²) per feature.
Instead, sort the feature once. A cumulative sum of one-hot labels gives left counts at every cut. Right counts are total minus left. All n - 1 cuts cost O(nC) after the sort.
Only cut between two different values. Use the midpoint as the threshold.

Solution

import numpy as np


def impurity(counts: np.ndarray, kind: str = "gini") -> np.ndarray:
    """counts: (..., C) class counts. Returns impurity per row. Empty rows give 0."""
    tot = counts.sum(-1, keepdims=True)
    p = np.divide(counts, tot, out=np.zeros_like(counts, dtype=float), where=tot > 0)
    if kind == "gini":
        return 1.0 - (p * p).sum(-1)
    logp = np.log2(p, out=np.zeros_like(p), where=p > 0)   # 0 log 0 = 0
    return -(p * logp).sum(-1)


def best_split(X: np.ndarray, y: np.ndarray, kind: str = "gini",
               min_leaf: int = 1) -> tuple[int, float, float]:
    """Find the (feature, threshold) that most reduces weighted impurity.

    X: (n, d) numeric. y: (n,) ints 0..C-1. Rule: go left if x <= threshold.
    Returns (feature, threshold, gain). feature is -1 if no valid split.
    """
    n, d = X.shape
    C = int(y.max()) + 1
    parent = impurity(np.bincount(y, minlength=C).astype(float), kind)
    best = (-1, np.nan, 0.0)
    for f in range(d):
        order = np.argsort(X[:, f], kind="stable")       # O(n log n)
        xs, ys = X[order, f], y[order]
        left = np.cumsum(np.eye(C)[ys], axis=0)[:-1]     # (n-1, C): counts left of cut i
        right = left[-1] + np.eye(C)[ys[-1]] - left      # total minus left
        n_left = np.arange(1, n)
        w_imp = (n_left * impurity(left, kind) + (n - n_left) * impurity(right, kind)) / n
        valid = xs[1:] > xs[:-1]                         # never cut between equal values
        valid &= (n_left >= min_leaf) & (n - n_left >= min_leaf)
        if not valid.any():
            continue
        w_imp = np.where(valid, w_imp, np.inf)
        i = int(w_imp.argmin())
        gain = parent - w_imp[i]
        if gain > best[2]:
            best = (f, (xs[i] + xs[i + 1]) / 2.0, float(gain))  # midpoint threshold
    return best


if __name__ == "__main__":
    X = np.array([[2.0, 7.0], [3.0, 1.0], [10.0, 5.0], [11.0, 3.0], [12.0, 9.0]])
    y = np.array([0, 0, 1, 1, 1])
    f, t, g = best_split(X, y)
    assert f == 0 and t == 6.5 and np.isclose(g, 0.48), (f, t, g)
    f2, t2, g2 = best_split(X, y, kind="entropy")
    assert f2 == 0 and np.isclose(g2, 0.97095, atol=1e-4)
    assert best_split(np.ones((4, 1)), np.array([0, 1, 0, 1]))[0] == -1
    print("ok", f, t, round(g, 3))

Line by line

impurity. It works on any leading shape, so one call scores every cut. np.divide(..., where=tot > 0) avoids dividing by zero.
np.log2(p, where=p > 0). It skips zeros. This applies the rule that 0 log 0 equals 0.
np.cumsum(np.eye(C)[ys], axis=0)[:-1]. Row i holds the class counts of the first i + 1 sorted rows. We drop the last row because a cut there leaves the right side empty.
valid = xs[1:] > xs[:-1]. You cannot split two equal values apart with a threshold. These cuts are masked out.
min_leaf. It blocks tiny children. Real trees use it to fight overfitting.
Midpoint threshold. It sits halfway between the two values. That generalizes better than either edge.

TimeO(d · n (log n + C))SpaceO(n C)

Follow-ups they will ask

Gini or entropy? They pick the same split most of the time. Gini skips the log, so it is a bit faster. That is why it is the sklearn default.
Regression trees? Use variance as the impurity. Keep cumulative sums of y and y². Then each cut's variance is O(1).
Categorical features? For binary targets, sort categories by their positive rate. Then scan like a numeric feature. This finds the best binary partition.
How does XGBoost or LightGBM do this faster? They bucket each feature into about 256 bins. Then they scan a histogram, not sorted rows.
Why is a greedy split not optimal? Finding the best full tree is NP-hard. Greedy picks the best split now and never looks ahead.

Common traps

Splitting between equal values. The threshold then cannot reproduce the split.
Comparing raw child impurities without weighting by size.
Taking log(0) in entropy.
Re-sorting inside the threshold loop. That brings back O(n² log n).

Say this out loud: “I sort each feature once. A cumulative sum of one-hot labels gives class counts on the left of every cut. Right counts are total minus left. So I score all cuts in one vectorized pass, and I skip cuts between equal values.”

7. ROC-AUC from Scores Medium

The prompt

“You have binary labels y (n,) and model scores s (n,). Compute ROC-AUC without sklearn. Scores can tie. Then also give me the ROC curve.”

What they are testing

Do you know AUC is the chance a random positive outranks a random negative?
Can you get it in O(n log n) with ranks, not O(n²) with pairs?
Do you count ties as one half?
Can you build the curve by sweeping the threshold and grouping tied scores?

Approach

Rank method. Rank all scores from 1 to n. Tied scores share their average rank.
Sum the ranks of the positives. Subtract n_pos(n_pos + 1)/2. That is the Mann-Whitney U. It counts the pairs where the positive wins, with ties as one half.
AUC = U / (n_pos · n_neg).
Sweep method. Sort scores high to low. Cumulative sums give TP and FP counts. Keep only the last row of each tied group. Then integrate with trapezoids.

Solution

import numpy as np


def rankdata_avg(a: np.ndarray) -> np.ndarray:
    """1-based ranks. Tied values share the average of their ranks."""
    order = np.argsort(a, kind="mergesort")
    s = a[order]
    # Start index of each run of equal values.
    starts = np.flatnonzero(np.r_[True, s[1:] != s[:-1]])
    ends = np.r_[starts[1:], len(s)]
    avg = (starts + ends + 1) / 2.0                 # mean of ranks starts+1 .. ends
    ranks = np.empty(len(a))
    ranks[order] = np.repeat(avg, ends - starts)
    return ranks


def roc_auc_rank(y: np.ndarray, s: np.ndarray) -> float:
    """AUC = P(score of random positive > random negative), ties count 1/2."""
    y = y.astype(bool)
    n_pos, n_neg = y.sum(), (~y).sum()
    if n_pos == 0 or n_neg == 0:
        raise ValueError("AUC needs both classes")
    r = rankdata_avg(s)
    u = r[y].sum() - n_pos * (n_pos + 1) / 2.0      # Mann-Whitney U
    return float(u / (n_pos * n_neg))


def roc_curve(y: np.ndarray, s: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Sweep the threshold from high to low. One point per distinct score."""
    order = np.argsort(-s, kind="mergesort")
    s, y = s[order], y[order].astype(float)
    last = np.r_[s[1:] != s[:-1], True]             # last index of each tied group
    tp = np.cumsum(y)[last]
    fp = np.cumsum(1.0 - y)[last]
    tpr = np.r_[0.0, tp / tp[-1]]
    fpr = np.r_[0.0, fp / fp[-1]]
    return fpr, tpr


def roc_auc_sweep(y: np.ndarray, s: np.ndarray) -> float:
    fpr, tpr = roc_curve(y, s)
    return float(np.sum(np.diff(fpr) * (tpr[1:] + tpr[:-1]) / 2.0))  # trapezoids


if __name__ == "__main__":
    y = np.array([0, 0, 1, 1, 0, 1])
    s = np.array([0.1, 0.4, 0.35, 0.8, 0.4, 0.4])
    brute = np.mean([(sp > sn) + 0.5 * (sp == sn) for sp in s[y == 1] for sn in s[y == 0]])
    assert np.isclose(roc_auc_rank(y, s), brute)
    assert np.isclose(roc_auc_sweep(y, s), brute)
    rng = np.random.default_rng(4)
    y2 = rng.integers(0, 2, 1000)
    s2 = np.round(rng.normal(size=1000) + y2, 1)    # rounding makes many ties
    assert np.isclose(roc_auc_rank(y2, s2), roc_auc_sweep(y2, s2))
    print("ok", round(brute, 4))

Line by line

starts. In sorted order, a new run starts where a value differs from the one before. flatnonzero turns that mask into indices.
(starts + ends + 1) / 2. The mean of ranks starts+1 through ends. This is the average rank for a tied run.
ranks[order] = np.repeat(...). It spreads each run's rank to its members and puts them back in the original order.
u = r[y].sum() - n_pos(n_pos+1)/2. If all positives were on the bottom, their rank sum would be exactly that subtracted term. What is left counts wins over negatives.
last = ... in roc_curve. All rows with the same score cross the threshold together. So we keep one point per group. The trapezoid between points then gives ties half credit.

TimeO(n log n)SpaceO(n)Brute forceO(n_pos · n_neg)

Follow-ups they will ask

Why do both methods agree? The trapezoid area under the step curve counts each pair once. A tied pair lands on a diagonal segment and gets half credit.
Is AUC good for heavy imbalance? It can look great while precision is poor. Report PR-AUC or precision at a fixed recall too.
Does AUC care about calibration? No. Any monotone change to the scores leaves AUC the same. Check calibration with log loss or a reliability plot.
Confidence interval? Bootstrap the rows and recompute. Or use the DeLong method for a closed form.
Too much data for one machine? Bucket scores into a histogram per class. Then merge counts across workers and sweep the buckets.

Common traps

Using plain argsort ranks. Ties then get arbitrary ranks and the answer moves with the sort order.
Adding a curve point inside a tied group. That gives a staircase that depends on row order.
Not handling a single-class input. AUC is undefined, so raise an error.
Sorting low to high in the sweep and getting 1 - AUC.

Say this out loud: “AUC is the chance a random positive scores above a random negative, with ties worth half. That equals the Mann-Whitney U over n_pos times n_neg. So I compute average ranks once and get it in n log n.”

8. Self-Attention Forward Pass Hard

The prompt

“Implement scaled dot-product attention in NumPy. Start with one head. Then make it multi-head with input X (B, T, D) and weights Wq, Wk, Wv, Wo, each (D, D). Add an option for a causal mask.”

What they are testing

Do you know the formula softmax(Q Kᵀ / √d) V and why the scale is there?
Can you split and merge heads with reshape and transpose correctly?
Do you mask before the softmax, with -inf?
Can you test that the mask really blocks the future?

Approach

Scores are Q Kᵀ / √d, shape (T, T). Row i says how much token i looks at each token.
For causal attention, set scores above the diagonal to -inf. After softmax their weight is exactly 0.
Softmax over the last axis, then multiply by V.
For multi-head, project once to (B, T, D). Reshape to (B, T, H, hd) and move heads forward to (B, H, T, hd). Run attention on all heads at once. Undo the move and apply Wo.

Solution

import numpy as np


def softmax(z: np.ndarray, axis: int = -1) -> np.ndarray:
    e = np.exp(z - z.max(axis=axis, keepdims=True))
    return e / e.sum(axis=axis, keepdims=True)


def attention(Q: np.ndarray, K: np.ndarray, V: np.ndarray,
              causal: bool = False) -> np.ndarray:
    """Scaled dot-product attention. Q, K, V: (..., T, d). Returns (..., T, d_v)."""
    d = Q.shape[-1]
    scores = Q @ np.swapaxes(K, -1, -2) / np.sqrt(d)      # (..., T, T)
    if causal:
        T = scores.shape[-1]
        mask = np.triu(np.ones((T, T), dtype=bool), k=1)  # True above the diagonal
        scores = np.where(mask, -np.inf, scores)          # block future positions
    return softmax(scores, axis=-1) @ V


def multi_head_attention(X: np.ndarray, Wq: np.ndarray, Wk: np.ndarray, Wv: np.ndarray,
                         Wo: np.ndarray, n_heads: int, causal: bool = False) -> np.ndarray:
    """X: (B, T, D). All W: (D, D). D must divide by n_heads. Returns (B, T, D)."""
    B, T, D = X.shape
    hd = D // n_heads

    def split(M: np.ndarray) -> np.ndarray:               # (B, T, D) -> (B, H, T, hd)
        return M.reshape(B, T, n_heads, hd).transpose(0, 2, 1, 3)

    Q, K, V = split(X @ Wq), split(X @ Wk), split(X @ Wv)
    out = attention(Q, K, V, causal=causal)               # (B, H, T, hd)
    out = out.transpose(0, 2, 1, 3).reshape(B, T, D)      # concat heads
    return out @ Wo


if __name__ == "__main__":
    rng = np.random.default_rng(5)
    B, T, D, H = 2, 4, 8, 2
    X = rng.normal(size=(B, T, D))
    W = [rng.normal(size=(D, D)) / np.sqrt(D) for _ in range(4)]
    out = multi_head_attention(X, *W, n_heads=H, causal=True)
    assert out.shape == (B, T, D)
    # Causal check: changing the last token must not change earlier outputs.
    X2 = X.copy()
    X2[:, -1] += 10.0
    out2 = multi_head_attention(X2, *W, n_heads=H, causal=True)
    assert np.allclose(out[:, :-1], out2[:, :-1]) and not np.allclose(out[:, -1], out2[:, -1])
    # One head equals plain attention on the projections.
    one = multi_head_attention(X, W[0], W[1], W[2], np.eye(D), n_heads=1)
    assert np.allclose(one, attention(X @ W[0], X @ W[1], X @ W[2]))
    print("ok", out.shape)

Line by line

np.swapaxes(K, -1, -2). It transposes only the last two axes. So the same code works for (T, d) and (B, H, T, d).
/ np.sqrt(d). Dot products of random d-dim vectors have variance about d. Dividing keeps the softmax out of its flat, saturated zone.
np.triu(..., k=1). True strictly above the diagonal. Those are the future positions.
np.where(mask, -np.inf, scores). Every row keeps its diagonal, so no row is all -inf. That means no nan from the softmax.
split. The reshape cuts D into H chunks of hd. The transpose brings heads next to batch so @ treats them as batch dims.
The tests. Changing the last token must not change earlier outputs. One head with Wo = I must match plain attention.

TimeO(B T² D + B T D²)SpaceO(B H T²)

Follow-ups they will ask

Why scale by √d? Without it, large scores push softmax to one-hot. Gradients then vanish for every other token.
Add a padding mask. Pass a (B, T) boolean of real tokens. Broadcast it to (B, 1, 1, T) and set padded keys to -inf. Watch for rows that are fully masked.
What is the memory problem? The T×T score matrix. FlashAttention tiles the work and uses an online softmax, so it never stores the full matrix.
What does a KV cache do? At decode time, keep past K and V. Each new token computes one query row against the cache. Cost per step drops from O(T²) to O(T).
Multi-query or grouped-query attention? Many query heads share fewer K and V heads. The KV cache shrinks by that ratio.

Common traps

Reshaping (B, T, D) straight to (B, H, T, hd). That mixes tokens across heads. Reshape to (B, T, H, hd) first, then transpose.
Masking after the softmax. The rows no longer sum to 1.
Using a large negative like -1e9 in float16. It overflows. Use -inf or the dtype's min.
Softmax over the wrong axis. It must be over keys, the last axis.

Say this out loud: “Scores are Q K transpose over root d, softmaxed over keys, times V. For causal I fill the upper triangle with minus infinity before the softmax. For heads I reshape to B, T, H, head-dim, then transpose heads forward and batch them.”

9. Two-Layer MLP with Backprop Hard

The prompt

“Build a two-layer network for classification. Input X (n, d), one hidden ReLU layer, softmax output over C classes, cross-entropy loss. Write the forward pass, the backward pass by hand, and a gradient check.”

What they are testing

Can you apply the chain rule through matmul, ReLU and softmax?
Do your gradient shapes match your parameter shapes?
Do you know a good init for ReLU?
Can you prove your gradients right with finite differences?

Approach

Forward: h = ReLU(X W1 + b1), then logits = h W2 + b2, then stable log-softmax and mean CE.
Backward, top down: dlogits = (softmax - onehot)/n. Then dW2 = hᵀ dlogits and db2 is the column sum.
Pass back: dh = dlogits W2ᵀ. ReLU lets the gradient through only where h_pre > 0.
Then dW1 = Xᵀ dh_pre and db1 is its column sum.
Grad check: nudge each weight by ±ε. Compare (L+ - L-)/2ε to the analytic gradient with relative error.

Solution

import numpy as np


def init_params(d_in: int, d_hidden: int, d_out: int, seed: int = 0) -> dict[str, np.ndarray]:
    rng = np.random.default_rng(seed)
    return {
        "W1": rng.normal(size=(d_in, d_hidden)) * np.sqrt(2.0 / d_in),  # He init for ReLU
        "b1": np.zeros(d_hidden),
        "W2": rng.normal(size=(d_hidden, d_out)) * np.sqrt(1.0 / d_hidden),
        "b2": np.zeros(d_out),
    }


def forward_backward(p: dict[str, np.ndarray], X: np.ndarray,
                     y: np.ndarray) -> tuple[float, dict[str, np.ndarray]]:
    """X: (n, d_in). y: (n,) class ids. Returns (mean CE loss, grads)."""
    n = X.shape[0]
    # Forward
    h_pre = X @ p["W1"] + p["b1"]                 # (n, H)
    h = np.maximum(h_pre, 0.0)                    # ReLU
    logits = h @ p["W2"] + p["b2"]                # (n, C)
    z = logits - logits.max(1, keepdims=True)
    logp = z - np.log(np.exp(z).sum(1, keepdims=True))
    loss = -logp[np.arange(n), y].mean()
    # Backward
    dlogits = np.exp(logp)
    dlogits[np.arange(n), y] -= 1.0
    dlogits /= n                                  # (n, C)
    grads = {"W2": h.T @ dlogits, "b2": dlogits.sum(0)}
    dh = dlogits @ p["W2"].T                      # (n, H)
    dh_pre = dh * (h_pre > 0)                     # ReLU passes grad only where active
    grads["W1"] = X.T @ dh_pre
    grads["b1"] = dh_pre.sum(0)
    return float(loss), grads


def grad_check(p: dict[str, np.ndarray], X: np.ndarray, y: np.ndarray,
               eps: float = 1e-5) -> float:
    """Max relative error between analytic and centered-difference grads."""
    _, g = forward_backward(p, X, y)
    worst = 0.0
    for name, W in p.items():
        it = np.nditer(W, flags=["multi_index"])
        for _ in it:
            i = it.multi_index
            old = W[i]
            W[i] = old + eps
            lp, _ = forward_backward(p, X, y)
            W[i] = old - eps
            lm, _ = forward_backward(p, X, y)
            W[i] = old                            # always restore
            num = (lp - lm) / (2 * eps)
            rel = abs(num - g[name][i]) / max(1e-8, abs(num) + abs(g[name][i]))
            worst = max(worst, rel)
    return worst


def train(X: np.ndarray, y: np.ndarray, hidden: int = 32, lr: float = 0.5,
          epochs: int = 500) -> dict[str, np.ndarray]:
    p = init_params(X.shape[1], hidden, int(y.max()) + 1)
    for _ in range(epochs):
        _, g = forward_backward(p, X, y)
        for k in p:
            p[k] -= lr * g[k]
    return p


if __name__ == "__main__":
    rng = np.random.default_rng(6)
    X = rng.normal(size=(6, 3))
    y = np.array([0, 1, 2, 1, 0, 2])
    assert grad_check(init_params(3, 5, 3, seed=1), X, y) < 1e-6
    # XOR is not linearly separable. The MLP should solve it.
    Xx = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=float)
    yx = np.array([0, 1, 1, 0])
    p = train(Xx, yx, hidden=16, lr=0.5, epochs=2000)
    pred = (np.maximum(Xx @ p["W1"] + p["b1"], 0) @ p["W2"] + p["b2"]).argmax(1)
    assert pred.tolist() == yx.tolist(), pred
    print("ok")

Line by line

He init. sqrt(2 / d_in) keeps the activation scale steady through a ReLU layer. ReLU zeros half the inputs, so we double the variance.
Bias gradients use .sum(0). The bias was broadcast over n rows in the forward pass. So its gradient sums over those rows.
Shape rule. A weight's gradient has the weight's shape. h.T @ dlogits is (H, n) @ (n, C), which gives (H, C), the shape of W2.
dh * (h_pre > 0). The ReLU derivative is 1 where active and 0 elsewhere.
grad_check. It edits each weight in place and always restores it. Relative error under 1e-6 in float64 means the backward pass is right.
XOR test. A linear model cannot fit XOR. The hidden layer can. It is a fast sanity check.

Step timeO(n (d H + H C))SpaceO(n H + d H + H C)Grad checkO(params · step)

Follow-ups they will ask

Why centered differences? Their error is O(ε²). A one-sided difference has O(ε) error.
What if ReLU units die? A unit with negative input everywhere gets zero gradient forever. Use Leaky ReLU, a lower learning rate, or better init.
Add L2 weight decay. Add λ W to each weight gradient. Leave the biases alone.
Add momentum or Adam. Keep a running average of gradients per parameter. Adam also keeps the running average of squared gradients and divides by its square root.
Why not init all weights to zero? Every hidden unit would get the same gradient. They would stay equal forever. Random init breaks the symmetry.

Common traps

Forgetting the / n in dlogits. Gradients then scale with batch size.
Using the post-ReLU h for the mask instead of h_pre. Here they agree, but not for other activations.
Grad-checking in float32. Round-off swamps the signal. Use float64.
Grad-checking at a kink. A ReLU input right at 0 gives a false alarm. Use a small random batch.

Say this out loud: “I go top down. The output gradient is softmax minus one-hot over n. Each weight gradient is the layer input transposed times the upstream gradient. ReLU masks the gradient where the pre-activation was negative. Then I verify everything with centered differences.”

10. K-Fold Cross-Validation Easy

The prompt

“Write k-fold cross-validation. Input is X (n, d), y (n,), a function that fits and predicts, and a metric. Shuffle the rows. Add a stratified option. Return the mean and standard deviation of the metric across folds.”

What they are testing

Do the folds cover every row exactly once?
Do you shuffle with a seed?
Can you keep class ratios the same in each fold?
Do you know about leakage, like fitting a scaler on all the data?

Approach

Plain k-fold. Permute 0..n-1 and cut it into k nearly equal parts with array_split.
Stratified. For each class, shuffle its rows. Deal them out to folds one at a time, like cards. Each fold then gets its fair share of every class.
For each fold, train on the rest and score on the fold.
Return the mean and the sample standard deviation of the k scores.

Solution

from typing import Callable

import numpy as np


def kfold_indices(y: np.ndarray, k: int = 5, stratified: bool = False,
                  seed: int = 0) -> list[np.ndarray]:
    """Split row ids 0..n-1 into k disjoint test folds. Shuffled."""
    rng = np.random.default_rng(seed)
    n = len(y)
    if k < 2 or k > n:
        raise ValueError("need 2 <= k <= n")
    if not stratified:
        return np.array_split(rng.permutation(n), k)
    folds: list[list[int]] = [[] for _ in range(k)]
    offset = 0
    for c in np.unique(y):
        ids = rng.permutation(np.flatnonzero(y == c))
        for j, i in enumerate(ids):
            folds[(j + offset) % k].append(int(i))   # deal like cards
        offset += len(ids)                           # next class starts where this one stopped
    return [np.array(sorted(f)) for f in folds]


def cross_validate(X: np.ndarray, y: np.ndarray,
                   fit_predict: Callable[[np.ndarray, np.ndarray, np.ndarray], np.ndarray],
                   metric: Callable[[np.ndarray, np.ndarray], float],
                   k: int = 5, stratified: bool = False, seed: int = 0) -> tuple[float, float]:
    """fit_predict(X_tr, y_tr, X_te) -> predictions. Returns (mean, std) of metric."""
    scores = []
    all_ids = np.arange(len(y))
    for test in kfold_indices(y, k, stratified, seed):
        train = np.setdiff1d(all_ids, test, assume_unique=True)
        # Fit every preprocessing step on train only, inside fit_predict.
        pred = fit_predict(X[train], y[train], X[test])
        scores.append(metric(y[test], pred))
    s = np.array(scores)
    return float(s.mean()), float(s.std(ddof=1))


if __name__ == "__main__":
    rng = np.random.default_rng(7)
    y = np.array([0] * 90 + [1] * 10)
    folds = kfold_indices(y, 5, stratified=True)
    assert sorted(np.concatenate(folds).tolist()) == list(range(100))
    assert all(y[f].sum() == 2 for f in folds)        # 10 positives over 5 folds

    X = rng.normal(size=(100, 2)) + y[:, None] * 3.0

    def nearest_mean(Xtr: np.ndarray, ytr: np.ndarray, Xte: np.ndarray) -> np.ndarray:
        mu = np.stack([Xtr[ytr == c].mean(0) for c in (0, 1)])
        return ((Xte[:, None, :] - mu[None]) ** 2).sum(-1).argmin(1)

    acc = lambda t, p: float((t == p).mean())
    m, s = cross_validate(X, y, nearest_mean, acc, k=5, stratified=True)
    assert m > 0.9
    print("ok", round(m, 3), round(s, 3))

Line by line

np.array_split. It allows parts of unequal size. Fold sizes differ by at most one.
folds[(j + offset) % k]. The round-robin deal. offset carries over between classes, so small folds do not always get the extra row.
np.setdiff1d(..., assume_unique=True). Train rows are all rows minus the test fold. The flag skips a needless sort check.
The fit_predict callback. Any scaling or feature selection must happen inside it, on train rows only. That keeps the test fold clean.
ddof=1. The sample standard deviation. With only k scores, the plain version runs low.

Split timeO(n log n)Totalk × fit costSpaceO(n)

Follow-ups they will ask

Time series? Never shuffle. Use forward-chaining splits. Train on the past and test on the next block. Leave a gap to stop leakage.
Many rows per user? Use group k-fold. All rows of one user go to one fold. Otherwise the model memorizes users.
Tuning hyperparameters too? Use nested CV. The inner loop picks settings. The outer loop scores the whole process.
Is the std a real confidence interval? No. The folds share training data, so scores are correlated. Treat it as a rough spread.
How to choose k? 5 or 10 is the norm. Larger k lowers bias but costs more fits and raises variance.

Common traps

Fitting the scaler, imputer or feature selector on all rows before splitting.
Stratifying by sorting labels without shuffling inside each class.
Forgetting the seed, so results change on every run.
Shuffling time-ordered data.

Say this out loud: “I shuffle the indices with a seed and split them into k folds. For the stratified version I deal each class out round-robin. All preprocessing is fit inside the loop on train rows only, so nothing leaks.”

Recap

Write the shape of every array in a comment. Most bugs are shape bugs.
Subtract the max before exp. Use log-sum-exp and logaddexp, never log(softmax).
The gradient of softmax or sigmoid with cross-entropy is always prediction minus target.
Pairwise distances are ||a||² + ||b||² - 2ab, clipped at zero.
Sort once and use cumulative sums. That trick powers tree splits, AUC and the ROC curve.
Test on a tiny case and grad-check with centered differences.

← 3 — Deep Learning and LLMs 5 — ML System Design →

ML Coding from Scratch

Contents

How ML coding rounds work

The NumPy you must know cold

The ten problems

1. Linear Regression by Gradient Descent Easy

The prompt

What they are testing

Approach

Solution

Line by line

Follow-ups they will ask

Common traps

2. Logistic Regression Easy

The prompt

What they are testing

Approach

Solution

Line by line

Follow-ups they will ask

Common traps

3. K-means Medium

The prompt

What they are testing

Approach

Solution

Line by line

Follow-ups they will ask

Common traps

4. K-Nearest Neighbours Medium

The prompt

What they are testing

Approach

Solution

Line by line

Follow-ups they will ask

Common traps

5. Stable Softmax and Cross-Entropy Easy

The prompt

What they are testing

Approach

Solution

Line by line

Follow-ups they will ask

Common traps

6. Best Decision-Tree Split Medium

The prompt

What they are testing

Approach

Solution

Line by line

Follow-ups they will ask

Common traps

7. ROC-AUC from Scores Medium

The prompt

What they are testing

Approach

Solution

Line by line

Follow-ups they will ask

Common traps

8. Self-Attention Forward Pass Hard

The prompt

What they are testing

Approach

Solution

Line by line

Follow-ups they will ask

Common traps

9. Two-Layer MLP with Backprop Hard

The prompt

What they are testing

Approach

Solution

Line by line

Follow-ups they will ask

Common traps

10. K-Fold Cross-Validation Easy

The prompt

What they are testing