def learnable_p_choose(
    energy,
    noise_mean: float = 0.0,
    noise_var: float = 0.0,
    training: bool = True
):
    """
    Calculating step wise prob for reading and writing
    1 to read, 0 to write
    energy: bsz, tgt_len, src_len
    """

    noise = 0
    if training:
        # add noise here to encourage discretness
        noise = (
            torch.normal(noise_mean, noise_var, energy.size())
            .type_as(energy)
            .to(energy.device)
        )

    p_choose = torch.sigmoid(energy + noise)

    # p_choose: bsz * self.num_heads, tgt_len, src_len
    return p_choose
Calculating step wise prob for reading and writing
1 to read, 0 to write
energy: bsz, tgt_len, src_len
examples/simultaneous_translation/utils/p_choose_strategy.py