def monotonic_attention_process_train(
        self,
        query: Optional[Tensor],
        key: Optional[Tensor],
        key_padding_mask: Optional[Tensor] = None,
    ):
        """
        Calculating monotonic attention process for training
        Including:
            stepwise probability: p_choose
            expected hard alignment: alpha
            expected soft attention: beta
        """
        assert query is not None
        assert key is not None

        # 1. compute stepwise probability
        p_choose = self.p_choose_from_qk(query, key, key_padding_mask)

        # 2. compute expected_alignment
        alpha = expected_alignment_from_p_choose(
            p_choose,
            key_padding_mask,
            eps=self.eps,
        )

        if self.mass_preservation:
            alpha = mass_preservation(
                alpha, key_padding_mask
            )

        # 3. compute expected soft attention (soft aligned model only)
        if self.soft_attention:
            soft_energy = self.energy_from_qk(
                query,
                key,
                "soft",
                key_padding_mask=None,
            )

            beta = expected_soft_attention(
                alpha,
                soft_energy,
                padding_mask=key_padding_mask,
                chunk_size=self.chunk_size,
                eps=self.eps,
            )
        else:
            beta = alpha
            soft_energy = alpha

        return p_choose, alpha, beta, soft_energy

        Calculating monotonic attention process for training
        Including:
            stepwise probability: p_choose
            expected hard alignment: alpha
            expected soft attention: beta
        
understand_icl/fairseq/examples/simultaneous_translation/modules/monotonic_multihead_attention.py