Here’s a thread which explains how the dimension orientation is handled in this assignment.
The tf.data.Dataset class does not change the orientation of the input data: it requires that the input have the samples dimension as the first dimension and then it subdivides along that dimension. The point about TF assuming samples are the first dimension is also covered in the thread I linked above.
It is generally a good idea to add the logic to process data in minibatches. The Dataset class in TF is commonly used for that. If you’re going to do that at all, then you do it for all your input datasets, because that’s the way you wrote the code.