src/transformers/models/esm/modeling_tf_esm.py/apply_rotary_pos_emb
def apply_rotary_pos_emb(x, cos, sin):
    cos = cos[:, :, : tf.shape(x)[-2], :]
    sin = sin[:, :, : tf.shape(x)[-2], :]

    return (x * cos) + (rotate_half(x) * sin)