src/transformers/models/distilbert/modeling_flax_distilbert.py/get_angles
def get_angles(pos, i, d_model):
    angle_rates = 1 / np.power(10000, (2 * (i // 2)) / np.float32(d_model))
    return pos * angle_rates