def _compute_loss(
        student_input_chunk,
        student_weight,
        teacher_input_chunk,
        teacher_weight,
        target_chunk,
        student_bias=None,
        teacher_bias=None,
        distillation_loss_fn=None,
        full_target=None,
        ignore_index=-100,
        weight_hard_loss=0.5,
        weight_soft_loss=0.5,
        compute_ce_loss=True,
        temperature=1,
        **loss_kwargs,
    ):
        """
        Compute the total loss for a chunk of input and target, while using an knowledge distillation loss function.
        Args:
            distillation_loss_fn (callable): Loss function to compute the loss on a chunk of input/target.
            student_input_chunk (torch.Tensor): Chunk of input tensor. Shape: (chunk_size, student_hidden_size).
            student_weight (torch.Tensor): Weight tensor. Shape: (vocab_size, student_hidden_size).
            teacher_input_chunk (torch.Tensor): Chunk of input tensor. Shape: (chunk_size, teacher_hidden_size).
            teacher_weight (torch.Tensor): Weight tensor. Shape: (vocab_size, teacher_hidden_size).
            target_chunk (torch.Tensor): Chunk of target tensor. Shape: (chunk_size,).
            student_bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,).
            teacher_bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,).
            full_target (torch.Tensor): Full target tensor. Shape: (batch_size * sequence_length,).
            ignore_index (int): Index to ignore for loss computation.
            weight_hard_loss (float): Weight for hard loss.
            weight_soft_loss (float): Weight for soft loss.
            compute_ce_loss (bool): Whether to compute CE loss.
            temperature (float): Temperature to control the input probability distribution. Default: `1.0` (i.e. no scale)
            loss_kwargs (dict): Additional arguments for the loss function.
        """
        (
            student_logits_chunk,
            teacher_logits_chunk,
            hard_loss,
        ) = LigerFusedLinearDistillationBase.chunk_forward(
            student_input_chunk,
            student_weight,
            teacher_input_chunk,
            teacher_weight,
            target_chunk,
            student_bias=student_bias,
            teacher_bias=teacher_bias,
            ignore_index=ignore_index,
            compute_ce_loss=compute_ce_loss,
        )

        student_logits_chunk /= temperature
        teacher_logits_chunk /= temperature

        # If the teacher and student token size is different, pad student logits to match the teacher's.
        # This only applies to cases where they share exactly the same vocab and tokenizer just
        # that teacher logit is padded for some training efficiency such as
        # https://huggingface.co/Qwen/Qwen1.5-72B-Chat/discussions/1#662883f568adf59b07b176d2
        teacher_vocab_size = teacher_weight.shape[0]
        student_vocab_size = student_weight.shape[0]
        if teacher_vocab_size > student_vocab_size:
            pad_size = teacher_vocab_size - student_vocab_size
            pad_tensor = torch.zeros(
                (*student_logits_chunk.shape[:-1], pad_size),
                dtype=student_logits_chunk.dtype,
                device=student_logits_chunk.device,
            )
            student_logits_chunk = torch.cat([student_logits_chunk, pad_tensor], dim=-1)

        hard_loss /= full_target.shape[0]

        soft_loss = distillation_loss_fn(student_logits_chunk, teacher_logits_chunk, **loss_kwargs)
        soft_loss /= full_target.shape[0]

        loss = weight_hard_loss * hard_loss + weight_soft_loss * soft_loss
        return loss, (soft_loss, hard_loss, student_logits_chunk, teacher_logits_chunk)

        Compute the total loss for a chunk of input and target, while using an knowledge distillation loss function.
        Args:
            distillation_loss_fn (callable): Loss function to compute the loss on a chunk of input/target.
            student_input_chunk (torch.Tensor): Chunk of input tensor. Shape: (chunk_size, student_hidden_size).
            student_weight (torch.Tensor): Weight tensor. Shape: (vocab_size, student_hidden_size).
            teacher_input_chunk (torch.Tensor): Chunk of input tensor. Shape: (chunk_size, teacher_hidden_size).
            teacher_weight (torch.Tensor): Weight tensor. Shape: (vocab_size, teacher_hidden_size).
            target_chunk (torch.Tensor): Chunk of target tensor. Shape: (chunk_size,).
            student_bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,).
            teacher_bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,).
            full_target (torch.Tensor): Full target tensor. Shape: (batch_size * sequence_length,).
            ignore_index (int): Index to ignore for loss computation.
            weight_hard_loss (float): Weight for hard loss.
            weight_soft_loss (float): Weight for soft loss.
            compute_ce_loss (bool): Whether to compute CE loss.
            temperature (float): Temperature to control the input probability distribution. Default: `1.0` (i.e. no scale)
            loss_kwargs (dict): Additional arguments for the loss function.
        
src/liger_kernel/chunked_loss/fused_linear_distillation.py