"""Distillation loss for matched tokenizers, contexts and generated positions.

Requires PyTorch. Inputs have shape [valid_positions, vocabulary_size].
This is a loss function, not a complete trainer or a tokenizer alignment method.
"""
import torch
import torch.nn.functional as F


def distillation_loss(student_logits, teacher_logprobs, temperature=2.0):
    if temperature <= 0:
        raise ValueError('temperature must be positive')
    if student_logits.ndim != 2 or student_logits.shape != teacher_logprobs.shape:
        raise ValueError('Inputs must have matching [valid_positions, vocabulary] shapes')
    if student_logits.shape[0] == 0:
        raise ValueError('At least one valid position is required')
    teacher = teacher_logprobs.detach().to(device=student_logits.device, dtype=torch.float32)
    log_p_teacher = F.log_softmax(teacher / temperature, dim=-1)
    log_p_student = F.log_softmax(student_logits.float() / temperature, dim=-1)
    return F.kl_div(
        log_p_student, log_p_teacher,
        reduction='batchmean', log_target=True,
    ) * temperature ** 2