Softmax Cross-Entropy

Task

Implement softmax_cross_entropy(logits: list[list[float]], targets: list[int]) -> float. Compute the mean multiclass cross-entropy loss from a batch of logits and integer class targets using a numerically stable log-sum-exp calculation.

Requirements

• logits is a non-empty rectangular batch with shape (N, C), where C is at least 2.
• targets contains one valid integer class index in [0, C) for each batch row.
• Compute normalization independently across the class dimension for each row.
• Use the log-sum-exp identity with the row maximum; do not exponentiate unshifted logits or compute log(softmax(logits)).
• For each row, subtract the target-class logit from that row's log-sum-exp.
• Return the arithmetic mean of the N per-example losses.

Example

logits = [[1.0, 2.0, 3.0], [1.0, 0.0, -1.0]]
targets = [2, 0]

softmax_cross_entropy(logits, targets)
# 0.4076059644