def _unflatten_and_set_shape(schema, ngram, fields_as_list):
    """Takes a flat list of fields produced by _flatten function and unflatten it back into an ngrams dictionary"""
    fields_as_namedtuple = make_namedtuple_tf_ngram(schema, ngram, *fields_as_list)

    # We change the key to str format here in order to be able to use ** later to expand the dictionary as kargs.
    fields_as_dict = {str(timestep): fields_as_namedtuple[timestep]._asdict() for timestep in fields_as_namedtuple}

    for timestep in fields_as_dict:
        _set_shape(schema, fields_as_dict[timestep])

    return make_namedtuple_tf_ngram(schema, ngram, **fields_as_dict)
Takes a flat list of fields produced by _flatten function and unflatten it back into an ngrams dictionary
petastorm/tf_utils.py