Cross-Entropy Without Overflow

Implement cross_entropy(logits, targets) that matches torch.nn.functional.cross_entropy with default arguments, without calling that function (or F.log_softmax).

Inputs:

• logits: shape (N, C), real-valued, unnormalized scores.
• targets: shape (N,), integer class indices in [0, C).

Output: a scalar tensor equal to the mean negative log-likelihood across the batch. Your implementation must handle very large logits (e.g. 1000.0) without overflowing to infinity.

Hints

• The standard trick for a stable log-softmax is to subtract logits.max(dim=-1, keepdim=True) before exponentiating. This preserves the softmax (it cancels) but keeps every exp argument <= 0.
• Equivalent and more compact: use torch.logsumexp to compute the log-normalizer in one fused, stable op.
• gather (or fancy indexing) picks out the log-probability of the correct class per row in a single call.