"""Derive and verify the softmax cross-entropy gradient on CPU."""

from __future__ import annotations

import torch


def stable_softmax(logits: torch.Tensor) -> torch.Tensor:
    """Return softmax probabilities without exponentiating the largest logit."""
    shifted = logits - logits.max(dim=-1, keepdim=True).values
    exponentials = shifted.exp()
    return exponentials / exponentials.sum(dim=-1, keepdim=True)


def one_hot(targets: torch.Tensor, classes: int, *, dtype: torch.dtype) -> torch.Tensor:
    """Turn integer class IDs shaped ``(B,)`` into distributions shaped ``(B, C)``."""
    return torch.nn.functional.one_hot(targets, num_classes=classes).to(dtype=dtype)


def softmax_cross_entropy(
    logits: torch.Tensor, targets: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Return mean loss, probabilities, and the exact mean-loss gradient.

    ``logits`` is ``(B, C)`` and integer ``targets`` is ``(B,)``.  The returned
    gradient has shape ``(B, C)`` and includes the factor ``1 / B`` introduced
    by taking a mean over the batch.
    """
    if logits.ndim != 2:
        raise ValueError("logits must have shape (batch, classes)")
    if targets.shape != logits.shape[:1] or targets.dtype != torch.long:
        raise ValueError("targets must be int64 class IDs shaped (batch,)")
    if not bool(((0 <= targets) & (targets < logits.shape[1])).all()):
        raise ValueError("every target must name an existing class")

    probabilities = stable_softmax(logits)
    labels = one_hot(targets, logits.shape[1], dtype=logits.dtype)
    # Computing logsumexp(logits) - chosen_logit as two float32 results can
    # lose their small difference when every logit is near 1e8.  log_softmax
    # performs the subtraction before the reduction and keeps that difference.
    log_probabilities = torch.log_softmax(logits, dim=-1)
    loss = -log_probabilities.gather(1, targets[:, None]).mean()
    gradient = (probabilities - labels) / logits.shape[0]
    return loss, probabilities, gradient


def main() -> None:
    logits = torch.tensor([[2.0, 1.0, 0.0]], requires_grad=True)
    targets = torch.tensor([1])
    loss, probabilities, manual_gradient = softmax_cross_entropy(logits, targets)
    (autograd_gradient,) = torch.autograd.grad(loss, logits)

    print("probabilities:", [round(value, 4) for value in probabilities[0].tolist()])
    print("loss:", round(loss.item(), 4))
    print("p - y:", [round(value, 4) for value in manual_gradient[0].tolist()])
    print("matches autograd:", torch.allclose(manual_gradient, autograd_gradient))
    print("gradient sum:", round(manual_gradient.sum().item(), 7))


if __name__ == "__main__":
    main()