def _compute_loss(
        input_chunk,
        weight,
        target_chunk,
        preference_labels_chunk,
        bias=None,
        preference_loss_fn=None,
        full_target=None,
        ignore_index=-100,
        use_ref_model=False,
        ref_input_chunk=None,
        ref_weight=None,
        ref_bias=None,
        average_log_prob=False,
        **loss_kwargs,
    ):
        """
        Compute the total loss for a chunk of input and target, while using an alignment/preference loss function.
        Args:
            preference_loss_fn (callable): Loss function to compute the loss on a chunk of input/target.
            input_chunk (torch.Tensor): Chunk of input tensor. Shape: (2 * chunk_size, sequence_length, hidden_size).
            weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size).
            target_chunk (torch.Tensor): Chunk of target tensor. Shape: (2 * chunk_size, sequence_length).
            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.
            use_ref_model (bool): Whether to use a reference model for the alignment loss.
            ref_weight (torch.Tensor): Reference weight tensor. Shape: (vocab_size, hidden_size).
            ref_bias (torch.Tensor, optional): Reference bias tensor. Shape: (vocab_size,).
            average_log_prob (bool): Whether to average the log probability per non-masked token.
            loss_kwargs (dict): Additional arguments for the loss function.
        """
        (
            log_prob_chunk,
            chosen_logps_sum,
            rejected_logps_sum,
            chosen_logits_sum,
            rejected_logits_sum,
        ) = LigerFusedLinearUnpairedPreferenceBase.chunk_forward(
            input_chunk,
            weight,
            target_chunk,
            preference_labels_chunk,
            bias=bias,
            ignore_index=ignore_index,
            average_log_prob=average_log_prob,
        )

        if use_ref_model:
            with torch.no_grad():
                (
                    ref_log_prob_chunk,
                    _,
                    _,
                    _,
                    _,
                ) = LigerFusedLinearUnpairedPreferenceBase.chunk_forward(
                    ref_input_chunk,
                    ref_weight,
                    target_chunk,
                    preference_labels_chunk,
                    ref_bias,
                    ignore_index=ignore_index,
                    average_log_prob=average_log_prob,
                )
            loss_kwargs["ref_log_prob_chunk"] = ref_log_prob_chunk

        preference_loss_outputs = preference_loss_fn(
            log_prob_chunk, preference_labels_chunk, full_target, **loss_kwargs
        )
        if isinstance(preference_loss_outputs, tuple):
            preference_loss_chunk, *aux_outputs = preference_loss_outputs
        else:
            preference_loss_chunk, aux_outputs = preference_loss_outputs, []

        return_vars = (
            chosen_logps_sum,
            rejected_logps_sum,
            chosen_logits_sum,
            rejected_logits_sum,
        )

        return preference_loss_chunk, (*return_vars, *aux_outputs)

        Compute the total loss for a chunk of input and target, while using an alignment/preference loss function.
        Args:
            preference_loss_fn (callable): Loss function to compute the loss on a chunk of input/target.
            input_chunk (torch.Tensor): Chunk of input tensor. Shape: (2 * chunk_size, sequence_length, hidden_size).
            weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size).
            target_chunk (torch.Tensor): Chunk of target tensor. Shape: (2 * chunk_size, sequence_length).
            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.
            use_ref_model (bool): Whether to use a reference model for the alignment loss.
            ref_weight (torch.Tensor): Reference weight tensor. Shape: (vocab_size, hidden_size).
            ref_bias (torch.Tensor, optional): Reference bias tensor. Shape: (vocab_size,).
            average_log_prob (bool): Whether to average the log probability per non-masked token.
            loss_kwargs (dict): Additional arguments for the loss function.
        
src/liger_kernel/chunked_loss/fused_linear_unpaired_preference.py