Implement the forward and backward passes of softmax with cross-entropy from scratch, then write a training loop for a linear classifier.

import numpy as np

def softmax_ce(logits, y):
    z = logits - logits.max(1, keepdims=True)
    p = np.exp(z)
    p /= p.sum(1, keepdims=True)
    n = len(y)
    loss = -np.log(p[np.arange(n), y] + 1e-12).mean()
    grad = p.copy()
    grad[np.arange(n), y] -= 1
    return loss, grad / n

def train(X, y, k, lr=0.1, epochs=100):
    W = np.zeros((X.shape[1], k))
    b = np.zeros(k)
    for _ in range(epochs):
        loss, g = softmax_ce(X @ W + b, y)
        W -= lr * X.T @ g
        b -= lr * g.sum(0)
    return W, b

Subtract the row maximum before exponentiating, for numerical stability. The gradient with respect to the logits is probabilities minus the one-hot label, averaged over the batch; the chain rule gives X transposed times that for the weights. Interviewers look for the stability trick, the averaging, and a finite-difference gradient check.