import random
import unittest

from torch.utils.data import BatchSampler, DataLoader, IterableDataset

from accelerate import Accelerator
from accelerate.data_loader import (
    BatchSamplerShard,
    DataLoaderDispatcher,
    DataLoaderShard,
    IterableDatasetShard,
    SkipBatchSampler,
    SkipDataLoader,
    skip_first_batches,
)


class 	lowerCamelCase_     (   UpperCAmelCase_ ):






    '''simple docstring'''






    def __init__(	self , __lowercase=0.01 , __lowercase=1_000)       ->	Tuple:
        __UpperCamelCase     :Optional[int]			  =		p_stop
        __UpperCamelCase     :List[Any]			  =		max_length







    def __iter__(	self)       ->	Dict:
        __UpperCamelCase     :Union[str, Any]			  =		0
        __UpperCamelCase     :List[str]			  =		False
        while not stop and count < self.max_length:
            yield count
            count += 1
            __UpperCamelCase     :Optional[int]			  =		random.random() < self.p_stop


class 	lowerCamelCase_     (   unittest.TestCase ):






    '''simple docstring'''






    def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase=False , __lowercase=True)       ->	List[str]:
        __UpperCamelCase     :Optional[int]			  =		[
            BatchSamplerShard(SCREAMING_SNAKE_CASE_ , 2 , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)
            for i in range(2)
        ]
        __UpperCamelCase     :int			  =		[list(SCREAMING_SNAKE_CASE_) for batch_sampler_shard in batch_sampler_shards]
        if not split_batches:
            self.assertListEqual([len(SCREAMING_SNAKE_CASE_) for shard in batch_sampler_shards] , [len(SCREAMING_SNAKE_CASE_) for e in expected])
        self.assertListEqual(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_)







    def 			UpperCamelCase__   (	self)       ->	Optional[Any]:
        # Check the shards when the dataset is a round multiple of total batch size.
        __UpperCamelCase     :Optional[int]			  =		BatchSampler(range(24) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Union[str, Any]			  =		[
            [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
            [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :Optional[int]			  =		BatchSampler(range(24) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        # Expected shouldn't change
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_)

        # Check the shards when the dataset is a round multiple of batch size but not total batch size.
        __UpperCamelCase     :List[str]			  =		BatchSampler(range(21) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Any			  =		[
            [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
            [[3, 4, 5], [9, 10, 11], [15, 16, 17], [0, 1, 2]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :str			  =		BatchSampler(range(21) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :str			  =		[
            [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
            [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_)

        # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
        # num_processes batch.
        __UpperCamelCase     :Dict			  =		BatchSampler(range(22) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :List[Any]			  =		[
            [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
            [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 0, 1]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :List[Any]			  =		BatchSampler(range(22) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :List[str]			  =		[
            [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
            [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_)

        # Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
        # num_processes batch.
        __UpperCamelCase     :int			  =		BatchSampler(range(20) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :List[str]			  =		[
            [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 0]],
            [[3, 4, 5], [9, 10, 11], [15, 16, 17], [1, 2, 3]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :int			  =		BatchSampler(range(20) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :str			  =		[
            [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
            [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_)

        # Check the shards when the dataset is very small.
        __UpperCamelCase     :Tuple			  =		BatchSampler(range(2) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Dict			  =		[[[0, 1, 0]], [[1, 0, 1]]]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :int			  =		BatchSampler(range(2) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Any			  =		[[], []]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_)







    def 			UpperCamelCase__   (	self)       ->	List[str]:
        # Check the shards when the dataset is a round multiple of batch size.
        __UpperCamelCase     :int			  =		BatchSampler(range(24) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Optional[Any]			  =		[
            [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
            [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :List[str]			  =		BatchSampler(range(24) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        # Expected shouldn't change
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_)

        # Check the shards when the dataset is not a round multiple of batch size.
        __UpperCamelCase     :Dict			  =		BatchSampler(range(22) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :List[Any]			  =		[
            [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
            [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [0, 1]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :List[Any]			  =		BatchSampler(range(22) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :List[Any]			  =		[
            [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
            [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_)

        # Check the shards when the dataset is not a round multiple of batch size or num_processes.
        __UpperCamelCase     :List[str]			  =		BatchSampler(range(21) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Any			  =		[
            [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 0]],
            [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [1, 2]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :Any			  =		BatchSampler(range(21) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :List[Any]			  =		[
            [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
            [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_)

        # Check the shards when the dataset is very small.
        __UpperCamelCase     :int			  =		BatchSampler(range(2) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Optional[Any]			  =		[[[0, 1]], [[0, 1]]]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :Union[str, Any]			  =		BatchSampler(range(2) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Optional[Any]			  =		[[], []]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_)







    def 			UpperCamelCase__   (	self)       ->	List[Any]:
        # Check the shards when the dataset is a round multiple of total batch size.
        __UpperCamelCase     :Any			  =		BatchSampler(range(24) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :List[str]			  =		[
            [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
            [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :Union[str, Any]			  =		BatchSampler(range(24) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        # Expected shouldn't change
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)

        # Check the shards when the dataset is a round multiple of batch size but not total batch size.
        __UpperCamelCase     :List[str]			  =		BatchSampler(range(21) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Optional[int]			  =		[
            [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
            [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :Any			  =		BatchSampler(range(21) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Union[str, Any]			  =		[
            [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
            [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)

        # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
        # num_processes batch.
        __UpperCamelCase     :int			  =		BatchSampler(range(22) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :List[Any]			  =		[
            [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
            [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :Optional[Any]			  =		BatchSampler(range(22) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :List[str]			  =		[
            [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
            [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)

        # Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
        # num_processes batch.
        __UpperCamelCase     :int			  =		BatchSampler(range(20) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Union[str, Any]			  =		[
            [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19]],
            [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :Dict			  =		BatchSampler(range(20) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Optional[int]			  =		[
            [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
            [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)

        # Check the shards when the dataset is very small.
        __UpperCamelCase     :str			  =		BatchSampler(range(2) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Optional[Any]			  =		[[[0, 1]], []]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :Any			  =		BatchSampler(range(2) , batch_size=3 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :List[str]			  =		[[], []]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)







    def 			UpperCamelCase__   (	self)       ->	Optional[int]:
        # Check the shards when the dataset is a round multiple of batch size.
        __UpperCamelCase     :str			  =		BatchSampler(range(24) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :List[Any]			  =		[
            [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
            [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :int			  =		BatchSampler(range(24) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        # Expected shouldn't change
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)

        # Check the shards when the dataset is not a round multiple of batch size.
        __UpperCamelCase     :Tuple			  =		BatchSampler(range(22) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Any			  =		[
            [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
            [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :str			  =		BatchSampler(range(22) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Dict			  =		[
            [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
            [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)

        # Check the shards when the dataset is not a round multiple of batch size or num_processes.
        __UpperCamelCase     :Union[str, Any]			  =		BatchSampler(range(21) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Union[str, Any]			  =		[
            [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20]],
            [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :Tuple			  =		BatchSampler(range(21) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Dict			  =		[
            [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
            [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
        ]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)

        # Check the shards when the dataset is very small.
        __UpperCamelCase     :List[str]			  =		BatchSampler(range(2) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :List[Any]			  =		[[[0, 1]], []]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :int			  =		BatchSampler(range(2) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Tuple			  =		[[], []]
        self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_)







    def 			UpperCamelCase__   (	self)       ->	str:
        __UpperCamelCase     :Optional[Any]			  =		[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
        __UpperCamelCase     :Union[str, Any]			  =		[BatchSamplerShard(SCREAMING_SNAKE_CASE_ , 2 , SCREAMING_SNAKE_CASE_ , even_batches=SCREAMING_SNAKE_CASE_) for i in range(2)]

        self.assertEqual(len(batch_sampler_shards[0]) , 3)
        self.assertEqual(len(batch_sampler_shards[1]) , 2)

        self.assertListEqual(list(batch_sampler_shards[0]) , [[0, 1, 2], [5, 6, 7, 8], [12, 13]])
        self.assertListEqual(list(batch_sampler_shards[1]) , [[3, 4], [9, 10, 11]])







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase=False , __lowercase=2 , __lowercase=False)       ->	Any:
        random.seed(SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :str			  =		list(SCREAMING_SNAKE_CASE_)

        __UpperCamelCase     :int			  =		[
            IterableDatasetShard(
                SCREAMING_SNAKE_CASE_ , batch_size=SCREAMING_SNAKE_CASE_ , drop_last=SCREAMING_SNAKE_CASE_ , num_processes=SCREAMING_SNAKE_CASE_ , process_index=SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_ , )
            for i in range(SCREAMING_SNAKE_CASE_)
        ]
        __UpperCamelCase     :int			  =		[]
        for iterable_dataset_shard in iterable_dataset_shards:
            # Since our random iterable dataset will be... random... we need to use a seed to get reproducible results.
            random.seed(SCREAMING_SNAKE_CASE_)
            iterable_dataset_lists.append(list(SCREAMING_SNAKE_CASE_))

        __UpperCamelCase     :Optional[int]			  =		batch_size // num_processes if split_batches else batch_size
        # All iterable dataset shard should have the same length, a round multiple of shard_batch_size
        __UpperCamelCase     :Dict			  =		iterable_dataset_lists[0]
        for l in iterable_dataset_lists[1:]:
            self.assertEqual(len(SCREAMING_SNAKE_CASE_) , len(SCREAMING_SNAKE_CASE_))
            self.assertTrue(len(SCREAMING_SNAKE_CASE_) % shard_batch_size == 0)

        __UpperCamelCase     :Union[str, Any]			  =		[]
        for idx in range(0 , len(SCREAMING_SNAKE_CASE_) , SCREAMING_SNAKE_CASE_):
            for l in iterable_dataset_lists:
                observed += l[idx : idx + shard_batch_size]

        if not drop_last:
            while len(SCREAMING_SNAKE_CASE_) < len(SCREAMING_SNAKE_CASE_):
                reference += reference
        self.assertListEqual(SCREAMING_SNAKE_CASE_ , reference[: len(SCREAMING_SNAKE_CASE_)])







    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        __UpperCamelCase     :Tuple			  =		42
        __UpperCamelCase     :int			  =		RandomIterableDataset()

        self.check_iterable_dataset_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_)
        self.check_iterable_dataset_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_)
        self.check_iterable_dataset_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_)
        self.check_iterable_dataset_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_)

        # Edge case with a very small dataset
        __UpperCamelCase     :Union[str, Any]			  =		RandomIterableDataset(max_length=2)

        self.check_iterable_dataset_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_)
        self.check_iterable_dataset_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_)
        self.check_iterable_dataset_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_)
        self.check_iterable_dataset_shards(SCREAMING_SNAKE_CASE_ , SCREAMING_SNAKE_CASE_ , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_ , split_batches=SCREAMING_SNAKE_CASE_)







    def 			UpperCamelCase__   (	self)       ->	List[Any]:
        __UpperCamelCase     :Dict			  =		BatchSampler(range(16) , batch_size=4 , drop_last=SCREAMING_SNAKE_CASE_)
        __UpperCamelCase     :Tuple			  =		SkipBatchSampler(SCREAMING_SNAKE_CASE_ , 2)
        self.assertListEqual(list(SCREAMING_SNAKE_CASE_) , [[8, 9, 10, 11], [12, 13, 14, 15]])







    def 			UpperCamelCase__   (	self)       ->	Any:
        __UpperCamelCase     :List[str]			  =		SkipDataLoader(list(range(16)) , batch_size=4 , skip_batches=2)
        self.assertListEqual([t.tolist() for t in dataloader] , [[8, 9, 10, 11], [12, 13, 14, 15]])







    def 			UpperCamelCase__   (	self)       ->	Optional[int]:
        __UpperCamelCase     :Union[str, Any]			  =		DataLoader(list(range(16)) , batch_size=4)
        __UpperCamelCase     :Any			  =		skip_first_batches(SCREAMING_SNAKE_CASE_ , num_batches=2)
        self.assertListEqual([t.tolist() for t in new_dataloader] , [[8, 9, 10, 11], [12, 13, 14, 15]])







    def 			UpperCamelCase__   (	self)       ->	Optional[int]:
        __UpperCamelCase     :str			  =		DataLoaderShard(list(range(16)) , batch_size=4)
        for idx, _ in enumerate(SCREAMING_SNAKE_CASE_):
            self.assertEqual(dataloader.end_of_dataloader , idx == 3)

        # Test it also works on the second iteration
        for idx, _ in enumerate(SCREAMING_SNAKE_CASE_):
            self.assertEqual(dataloader.end_of_dataloader , idx == 3)







    def 			UpperCamelCase__   (	self)       ->	Any:
        Accelerator()
        __UpperCamelCase     :Any			  =		DataLoaderDispatcher(range(16) , batch_size=4)
        for idx, _ in enumerate(SCREAMING_SNAKE_CASE_):
            self.assertEqual(dataloader.end_of_dataloader , idx == 3)

        # Test it also works on the second iteration
        for idx, _ in enumerate(SCREAMING_SNAKE_CASE_):
            self.assertEqual(dataloader.end_of_dataloader , idx == 3)







import unittest

import numpy as np
import torch

from .utils_summarization import build_mask, compute_token_type_ids, process_story, truncate_or_pad
class  lowercase__(				unittest.TestCase							):



	"""simple docstring"""




	def 						_lowercase  (       self  :			List[str]       )       ->     List[Any]:
					lowercase_       =					1_0
	def 						_lowercase  (       self  :			int       )       ->     List[str]:
					lowercase_       =					[1, 2, 3, 4]
					lowercase_       =					[1, 2, 3, 4, 0, 0, 0, 0, 0, 0]
					self.assertEqual(truncate_or_pad(SCREAMING_SNAKE_CASE_       ,							self.block_size       ,							0       )       ,							SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			int       )       ->     Optional[Any]:
					lowercase_       =					[1, 2, 3, 4, 5, 6, 7, 8, 9, 1_0]
					lowercase_       =					[1, 2, 3, 4, 5, 6, 7, 8, 9, 1_0]
					self.assertEqual(truncate_or_pad(SCREAMING_SNAKE_CASE_       ,							self.block_size       ,							0       )       ,							SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			Union[str, Any]       )       ->     Optional[int]:
					lowercase_       =					[1, 2, 3, 4, 5, 6, 7, 8, 9, 1_0, 1_1, 1_2, 1_3]
					lowercase_       =					[1, 2, 3, 4, 5, 6, 7, 8, 9, 1_0]
					self.assertEqual(truncate_or_pad(SCREAMING_SNAKE_CASE_       ,							self.block_size       ,							0       )       ,							SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			Any       )       ->     List[Any]:
					lowercase_       =					'''It was the year of Our Lord one thousand seven hundred and
        seventy-five.\n\nSpiritual revelations were conceded to England at that
        favoured period, as at this.'''
					lowercase_  ,	lowercase_       =					process_story(SCREAMING_SNAKE_CASE_       )
					self.assertEqual(SCREAMING_SNAKE_CASE_       ,							[]       )
	def 						_lowercase  (       self  :			List[str]       )       ->     List[str]:
					lowercase_       =					''''''
					lowercase_  ,	lowercase_       =					process_story(SCREAMING_SNAKE_CASE_       )
					self.assertEqual(SCREAMING_SNAKE_CASE_       ,							[]       )
					self.assertEqual(SCREAMING_SNAKE_CASE_       ,							[]       )
	def 						_lowercase  (       self  :			Union[str, Any]       )       ->     Union[str, Any]:
					lowercase_       =					(
					    '''It was the year of Our Lord one thousand seven hundred and '''
					    '''seventy-five\n\nSpiritual revelations were conceded to England '''
					    '''at that favoured period, as at this.\n@highlight\n\nIt was the best of times'''
					)
					lowercase_  ,	lowercase_       =					process_story(SCREAMING_SNAKE_CASE_       )

					lowercase_       =					[
					    '''It was the year of Our Lord one thousand seven hundred and seventy-five.''',
					    '''Spiritual revelations were conceded to England at that favoured period, as at this.''',
					]
					self.assertEqual(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       )

					lowercase_       =					['''It was the best of times.''']
					self.assertEqual(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			Union[str, Any]       )       ->     Optional[Any]:
					lowercase_       =					torch.tensor([1, 2, 3, 4]       )
					lowercase_       =					torch.tensor([1, 1, 1, 1]       )
					np.testing.assert_array_equal(build_mask(SCREAMING_SNAKE_CASE_       ,							0       ).numpy()       ,							expected.numpy()       )
	def 						_lowercase  (       self  :			List[Any]       )       ->     Tuple:
					lowercase_       =					torch.tensor([1, 2, 3, 4, 2_3, 2_3, 2_3]       )
					lowercase_       =					torch.tensor([1, 1, 1, 1, 0, 0, 0]       )
					np.testing.assert_array_equal(build_mask(SCREAMING_SNAKE_CASE_       ,							2_3       ).numpy()       ,							expected.numpy()       )
	def 						_lowercase  (       self  :			int       )       ->     Dict:
					lowercase_       =					torch.tensor([8, 2, 3, 4, 1, 1, 1]       )
					lowercase_       =					torch.tensor([1, 1, 1, 1, 0, 0, 0]       )
					np.testing.assert_array_equal(build_mask(SCREAMING_SNAKE_CASE_       ,							1       ).numpy()       ,							expected.numpy()       )




	def 						_lowercase  (       self  :			List[str]       )       ->     Tuple:
					lowercase_       =					1_0_1
					lowercase_       =					torch.tensor([[1, 2, 3, 4, 5, 6], [1, 2, 3, 1_0_1, 5, 6], [1, 1_0_1, 3, 4, 1_0_1, 6]]       )
					lowercase_       =					torch.tensor([[1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 0, 0], [1, 0, 0, 0, 1, 1]]       )

					lowercase_       =					compute_token_type_ids(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       )
					np.testing.assert_array_equal(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       )




