def distillation_loss_fn(
        student_logits,
        teacher_logits,
    ):
        """
        Compute distillation loss.
        Args:
            student_logits (torch.Tensor): Raw (temperature-scaled) logits of student tokens. Shape: (batch_size * seq_len, vocab_size).
            teacher_logits (torch.Tensor): Raw (temperature-scaled) logits of teacher tokens. Shape: (batch_size * seq_len, vocab_size).
        Returns:
            torch.Tensor: Sum of distillation losses for the chunk. The class will handle
                converting this to mean loss by dividing by the full batch size * sequence length in _compute_loss.
        """
        raise NotImplementedError("Distillation loss function must be implemented.")

        Compute distillation loss.
        Args:
            student_logits (torch.Tensor): Raw (temperature-scaled) logits of student tokens. Shape: (batch_size * seq_len, vocab_size).
            teacher_logits (torch.Tensor): Raw (temperature-scaled) logits of teacher tokens. Shape: (batch_size * seq_len, vocab_size).
        Returns:
            torch.Tensor: Sum of distillation losses for the chunk. The class will handle
                converting this to mean loss by dividing by the full batch size * sequence length in _compute_loss.
        
src/liger_kernel/chunked_loss/fused_linear_distillation.py