keras/src/trainers/data_adapters/tf_dataset_adapter.py/TFDatasetAdapter/batch_size
class TFDatasetAdapter:    def batch_size(self):
        first_element_spec = tree.flatten(self._dataset.element_spec)[0]
        return first_element_spec.shape[0]