def alignment_loss(
        self,
        policy_chosen_logps: torch.FloatTensor,
        policy_rejected_logps: torch.FloatTensor,
        ref_chosen_logps: torch.FloatTensor,
        ref_rejected_logps: torch.FloatTensor,
        kl: torch.FloatTensor = None,
    ):
        """Compute KTO loss for a batch of policy log probabilities.
        Args:
            policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,)
            policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,)
            ref_chosen_logps: Log probabilities of the reference model for the chosen responses. Shape: (batch_size,)
            ref_rejected_logps: Log probabilities of the reference model for the rejected responses. Shape: (batch_size,)
        Returns:
            The losses tensor contains the KTO loss for each example in the batch.
        """
        if kl is None:
            kl = torch.zeros(1).to(policy_chosen_logps.device)

        # Chosen losses
        chosen_logratios = policy_chosen_logps - ref_chosen_logps
        if policy_chosen_logps.shape[0] != 0 or ref_chosen_logps.shape[0] != 0:
            # Eqn (7) of the KTO paper (https://huggingface.co/papers/2402.01306)
            chosen_losses = 1 - F.sigmoid(self.beta * (chosen_logratios - kl))

        else:
            # lists can't be empty -- if they are, then accelerate.gather will hang
            chosen_losses = torch.Tensor([]).to(policy_chosen_logps.device)

        # Rejected losses
        rejected_logratios = policy_rejected_logps - ref_rejected_logps
        if policy_rejected_logps.shape[0] != 0 or ref_rejected_logps.shape[0] != 0:
            rejected_losses = 1 - F.sigmoid(self.beta * (kl - rejected_logratios))
        else:
            # lists can't be empty -- if they are, then accelerate.gather will hang
            rejected_losses = torch.Tensor([]).to(policy_rejected_logps.device)

        losses = torch.cat(
            (chosen_losses, rejected_losses),
            0,
        )

        chosen_rewards = self.beta * chosen_logratios
        rejected_rewards = self.beta * rejected_logratios

        return losses, chosen_rewards.sum(), rejected_rewards.sum()
Compute KTO loss for a batch of policy log probabilities.
        Args:
            policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,)
            policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,)
            ref_chosen_logps: Log probabilities of the reference model for the chosen responses. Shape: (batch_size,)
            ref_rejected_logps: Log probabilities of the reference model for the rejected responses. Shape: (batch_size,)
        Returns:
            The losses tensor contains the KTO loss for each example in the batch.
        
test/chunked_loss/test_kto_loss.py