Next-token Cross-Entropy loss (reuses Part 1 loss)

The problem, from first principles

[01-output-projection]'s logits give a score for every vocabulary word, at EVERY position in the sequence. Training a language model means adjusting its weights so that position t's logits ASSIGN HIGH PROBABILITY to whatever token actually occurs at position t + 1, exactly the "predict the next token" framing that gives autoregressive language models their name. [02-deep-learning-core/03-losses/02-cross-entropy]'s cross_entropy_forward already implements the general "penalize low probability assigned to the correct class" loss; the only new work here is correctly SHIFTING logits and targets by one position before handing them to that already-built function.

From theory to code

Implement next_token_cross_entropy_loss(logits, token_ids): logits at positions 0 through seq_len - 2 predict targets at positions 1 through seq_len - 1 (token_ids shifted forward by one), flattened and passed to [02-cross-entropy]'s cross_entropy_forward.

Constraints

• Position t's logits are compared against token_ids[t + 1], the token that ACTUALLY comes next, never token_ids[t] itself.
• The LAST position's logits have no "next token" to predict (there's nothing after the end of the sequence) and must be EXCLUDED, leaving seq_len - 1 valid (prediction, target) pairs from a length-seq_len sequence.
• Any leading batch dimensions are flattened together with the sequence dimension before calling cross_entropy_forward (which expects a flat (n, vocab_size)/(n,) pair).

Hints

vocab_size = predicted_logits.shape[-1]
flat_logits = predicted_logits.reshape(-1, vocab_size)
flat_targets = targets.reshape(-1)
return cross_entropy_forward(flat_logits, flat_targets)