src/transformers/models/pegasus/modeling_flax_pegasus.py/create_sinusoidal_positions
def create_sinusoidal_positions(n_pos, dim, dtype):
    position_enc = np.array([[pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)] for pos in range(n_pos)])
    sentinel = dim // 2 + dim % 2
    out = np.zeros_like(position_enc)
    out[:, 0:sentinel] = np.sin(position_enc[:, 0::2])
    out[:, sentinel:] = np.cos(position_enc[:, 1::2])

    return jnp.array(out, dtype=dtype)