import math


def scale_add_pos_emb(input_emb, pos_enc):
    """
    Scale input embs and add position encoding.

    :param input_emb: [bsz, seq_len, dim]
    :param pos_enc: [bsz, pos_len, dim] (pos_len must be greater than seq_len)
    :return: [bsz, seq_len, dim]
    """
    input_emb *= math.sqrt(input_emb.size(-1))
    input_emb += pos_enc(input_emb)
    return input_emb