def index_to_config_coord(config, index):
    """ Given an index of the flattened array, returns the 2d index of where
        this corresponds to the configs
    """
    config_shapes = [_.numel() for _ in config]
    assert index < sum(config_shapes)

    for i, config_len in enumerate(config_shapes):
        if index > config_len - 1:
            index -= config_len
        else:
            return (i, index)