examples/research_projects/performer/modeling_flax_performer_utils.py/_invert_perm
def _invert_perm(perm):
    perm_inv = [0] * len(perm)
    for i, j in enumerate(perm):
        perm_inv[j] = i
    return tuple(perm_inv)