The gradient of the softmax + cross-entropy loss with respect to the logits is beautifully simple:
Implement loss_and_grad(logits, labels) for logits of shape (N, C) and integer labels, returning (loss, grad), where loss is the mean cross-entropy and grad has the same shape as logits. Make it numerically stable.
Then implement numerical_grad(f, x, eps=1e-6): the central-difference gradient of a scalar function f at array x, perturbing each element by ±eps. It's the standard way to check a hand-written gradient.