src/transformers/models/gptj/modeling_tf_gptj.py/apply_rotary_pos_emb
def apply_rotary_pos_emb(tensor: tf.Tensor, sincos: tf.Tensor) -> tf.Tensor:
    sin_pos, cos_pos = sincos
    sin_pos = tf.repeat(sin_pos[:, :, None, :], 2, 3)
    cos_pos = tf.repeat(cos_pos[:, :, None, :], 2, 3)
    return (tensor * cos_pos) + (rotate_every_two(tensor) * sin_pos)