class IterableDatasetShard(IterableDataset):
    def __init__(
        self,
        dataset: IterableDataset,
        batch_size: int = 1,
        drop_last: bool = False,
        num_processes: int = 1,
        process_index: int = 0,
        split_batches: bool = False,
        group_by_length_buffer: int = 1, ## how many batches to preload
        length_column_name: str = "", ## which column determines length
    ):
        if split_batches and batch_size > 1 and batch_size % num_processes != 0:
            raise ValueError(
                f"To use `IterableDatasetShard` in `split_batches` mode, the batch size ({batch_size}) "
                f"needs to be a round multiple of the number of processes ({num_processes})."
            )
        self.dataset = dataset
        self.batch_size = batch_size
        self.drop_last = drop_last
        self.num_processes = num_processes
        self.process_index = process_index
        self.split_batches = split_batches
        self.group_by_length_buffer = group_by_length_buffer
        self.length_column_name = length_column_name

    def set_epoch(self, epoch):
        self.epoch = epoch
        if hasattr(self.dataset, "set_epoch"):
            self.dataset.set_epoch(epoch)

    def __len__(self):
        # We will just raise the downstream error if the underlying dataset is not sized
        if self.drop_last:
            return (len(self.dataset) // (self.batch_size * self.num_processes)) * self.batch_size
        else:
            return (
                math.ceil(len(self.dataset) / (self.batch_size * self.num_processes))
                * self.batch_size
            )

    def __iter__(self):
        if (
            not hasattr(self.dataset, "set_epoch")
            and hasattr(self.dataset, "generator")
            and isinstance(self.dataset.generator, torch.Generator)
        ):
            self.dataset.generator.manual_seed(self.epoch)
        real_batch_size = (
            self.batch_size if self.split_batches else (self.batch_size * self.num_processes)
        ) * self.group_by_length_buffer
        process_batch_size = (
            (self.batch_size // self.num_processes) if self.split_batches else self.batch_size
        ) * self.group_by_length_buffer
        process_slice = range(
            self.process_index * process_batch_size, (self.process_index + 1) * process_batch_size
        )

        first_batch = None
        current_batch = []
        for element in self.dataset:
            current_batch.append(element)
            # Wait to have a full batch before yielding elements.
            if len(current_batch) == real_batch_size:
                process_current_batch = [current_batch[i] for i in process_slice]
                process_current_batch.sort(key=lambda x: len(x[self.length_column_name]))
                yield from process_current_batch
                if first_batch is None:
                    first_batch = current_batch.copy()
                current_batch = []

        # Finished if drop_last is True, otherwise complete the last batch with elements from the beginning.
        if not self.drop_last and len(current_batch) > 0:
            if first_batch is None:
                first_batch = current_batch.copy()
            while len(current_batch) < real_batch_size:
                current_batch += first_batch
            for i in process_slice:
                yield current_batch[i]