"""Optional PyTorch CPU loss/gradient checks; no model download or LLM training."""
import torch
import torch.nn.functional as F


def masked_losses(student_logits, teacher_logits, labels, temperature=3.):
    # Caller has already shifted causal logits/labels; all arrays share [B,L].
    if temperature <= 0:
        raise ValueError('temperature must be positive')
    if student_logits.shape != teacher_logits.shape or labels.shape != student_logits.shape[:2]:
        raise ValueError('expected aligned logits [B,L,V] and labels [B,L]')
    mask = labels.ne(-100)
    if not mask.any():
        raise ValueError('batch contains no target token')
    s = student_logits.float()[mask]
    t = teacher_logits.detach().float()[mask]
    kd = F.kl_div(F.log_softmax(s/temperature, -1), F.softmax(t/temperature, -1), reduction='batchmean')*temperature**2
    ce = F.cross_entropy(s, labels[mask])
    return kd, ce


def main():
    torch.manual_seed(66)
    student = torch.randn(2,4,5, requires_grad=True)
    teacher = torch.randn(2,4,5, requires_grad=True)
    labels = torch.tensor([[-100,1,2,3],[-100,2,-100,-100]])
    kd,ce = masked_losses(student,teacher,labels)
    mask=labels.ne(-100)
    ps=F.log_softmax(student[mask]/3, -1)
    pt=F.log_softmax(teacher.detach()[mask]/3, -1)
    manual=(pt.exp()*(pt-ps)).sum(-1).mean()*9
    assert torch.allclose(kd,manual,atol=1e-6)
    (.7*kd+.3*ce).backward()
    assert teacher.grad is None
    assert student.grad[~mask].abs().sum().item()==0
    before=student.detach().clone()
    changed=before.clone(); changed[~mask]=1000.
    assert torch.allclose(masked_losses(changed,teacher,labels)[0],kd)
    try:
        masked_losses(student,teacher,torch.full_like(labels,-100))
    except ValueError:
        pass
    else:
        raise AssertionError('empty target batch not rejected')
    print({'torch':torch.__version__,'device':'cpu','logits_shape':list(student.shape),'valid_tokens':int(mask.sum()),'kd':kd.item(),'ce':ce.item(),'checks':'PASS'})


if __name__=='__main__':
    main()