Lectura del modelo TensorFlow de los archivos parquet utilizando IODataset.

He estado intentando usar tfio.IODataset.from_parquet para entrenar un modelo con grandes archivos parquet. A continuación se muestra un ejemplo mínimo del procedimiento de carga de parquet que estoy utilizando:

pd.DataFrame({'a': [.1, .2], 'b': [.01, .02]}).toparquet('file.parquet')
ds = tfio.IODataset.fromparquet('file.parquet', columns=['a', 'b'])

for batch in ds.batch(5):
    print(batch)

El conjunto de datos batched es de tipo OrderedDict con claves a y b. Para entrenar mi modelo, me gustaría tener algo más parecido a un vector de características “densas”, en lugar de dos claves separadas en un ordereddict. ¿Cómo puedo convertir OrderedDict en una tupla densa?

Intento 1

Según este ejemplo, intenté lo siguiente para transformar el conjunto de datos en características “densas”:

def make_dense(features):
    features = tf.stack(list(features), axis=1)
    return features
ds = ds.map(make_dense)

Desafortunadamente, eso arroja errores. He intentado varias variaciones sobre este tema, incluyendo
1. cambiar axis=1 a axis=0
2. usando ds = ds.map(lambda *items: tf.stack(items)) en lugar de mi función make_dense.

Imagino que esta es una operación muy básica para IODataset; simplemente no sé cómo lograrlo.