def METHOD_NAME(
    p_choose: Tensor,
    padding_mask: Optional[Tensor] = None,
    eps: float = 1e-6
):
    """
    Calculating expected alignment for from stepwise probability
    Reference:
    Online and Linear-Time Attention by Enforcing Monotonic Alignments
    https://arxiv.org/pdf/1704.00784.pdf
    q_ij = (1 − p_{ij−1})q_{ij−1} + a+{i−1j}
    a_ij = p_ij q_ij
    Parallel solution:
    ai = p_i * cumprod(1 − pi) * cumsum(a_i / cumprod(1 − pi))
    ============================================================
    Expected input size
    p_choose: bsz, tgt_len, src_len
    """
    prob_check(p_choose)
    # p_choose: bsz, tgt_len, src_len
    bsz, tgt_len, src_len = p_choose.size()
    dtype = p_choose.dtype
    p_choose = p_choose.float()
    if padding_mask is not None:
        p_choose = p_choose.masked_fill(padding_mask.unsqueeze(1), 0.0)
    if p_choose.is_cuda:
        p_choose = p_choose.contiguous()
        from alignment_train_cuda_binding import alignment_train_cuda as alignment_train
    else:
        from alignment_train_cpu_binding import alignment_train_cpu as alignment_train
    alpha = p_choose.new_zeros([bsz, tgt_len, src_len])
    alignment_train(p_choose, alpha, eps)
    # Mix precision to prevent overflow for fp16
    alpha = alpha.type(dtype)
    prob_check(alpha)
    return alpha