def repair_ids(shared_state, original_id_to_label, label_to_id):
    """Given an optimizer state and some new ids for the parameters,
    creates a new state with the new ids linked to the same tensors.

    Useful when using an optimizer state used on some old model on a
    new one.

    """

    new_state = {}
    for param_id, state in shared_state.items():
        new_param_id = label_to_id[original_id_to_label[param_id]]
        new_state[new_param_id] = state
    return new_state