I have done the predictions on batched data and want to map the predictions back to its input.

Is there any TensorFlow documentation which

- Explains how to map batched predictions back to it’s inputs.
- Confirms that batching won’t change the sequence of batched element and can be mapped to it’s input through index.

Our example - We are using row_number (dataframe.index) as a solution to map it back to it’s input

row_num = tf.keras.Input(shape =(1, ) , dtype = “int32”)

model1= tf.keras.Model( inputs = {“inputs” = model.layer[0].input , “row_number” :row_num },

output = {“encoding” : model.layers[4].output ,“row_number” :row_num }

)

def df_to_ds( series , batch_size)

dataset = tf.data.Dataset.from_tensor_slices(dict(series))

dataset = dataset.map(lambda e: ({“inputs” : e[“inputs”] , “row_number” : e[“row_number”] }))

dataset = dataset.batch( batch_size)

returns dataset

train_ds = df_to_ds(dataframe, batch_size = 64)

pred = model1.predict(train_ds)

This pred output will have a row_number which will be joined back to the original Dataframe through row_num