[tf.data] Validate that all elements of a batch have the same shape.

Fixes #15152. Previously, it would be possible for elements with
different shapes but the same overall number of elements to be batched
together, which would give unexpected behavior.

PiperOrigin-RevId: 178127613