def __init__(self, dataset, class_weight=None, distribution=None):
        """Initialize the TFDatasetAdapter.

        Args:
            dataset: The input `tf.data.Dataset` instance.
            class_weight: A map where the keys are integer class ids and values
                are the class weights, e.g. `{0: 0.2, 1: 0.6, 2: 0.3}`.
            distribution: A `keras.distribution.Distribution` instance. Used to
                shard the input dataset into per worker/process dataset
                instance.
        """
        from keras.src.utils.module_utils import tensorflow as tf

        if not isinstance(
            dataset, (tf.data.Dataset, tf.distribute.DistributedDataset)
        ):
            raise ValueError(
                "Expected argument `dataset` to be a tf.data.Dataset. "
                f"Received: {dataset}"
            )
        if class_weight is not None:
            dataset = dataset.map(
                make_class_weight_map_fn(class_weight)
            ).prefetch(tf.data.AUTOTUNE)
        if distribution is not None:
            dataset = distribution.distribute_dataset(dataset)
        self._dataset = dataset
Initialize the TFDatasetAdapter.

        Args:
            dataset: The input `tf.data.Dataset` instance.
            class_weight: A map where the keys are integer class ids and values
                are the class weights, e.g. `{0: 0.2, 1: 0.6, 2: 0.3}`.
            distribution: A `keras.distribution.Distribution` instance. Used to
                shard the input dataset into per worker/process dataset
                instance.
        
keras/src/trainers/data_adapters/tf_dataset_adapter.py