How to map predictions on batched dataset back to it's original input

General Discussion
#1

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

Is there any TensorFlow documentation which

  1. Explains how to map batched predictions back to it’s inputs.
  2. 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