keras/src/distribution/distribution_lib.py/Distribution/distribute_dataset
class Distribution:    def distribute_dataset(self, dataset):
        """Create a distributed dataset instance from the original user dataset.

        Args:
            dataset: the original global dataset instance. Only
            `tf.data.Dataset` is supported at the moment.

        Returns:
            a sharded `tf.data.Dataset` instance, which will produce data for
            the current local worker/process.
        """
        raise NotImplementedError()