


"""simple docstring"""



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 					__magic_name__		(			__UpperCAmelCase     ):






       def __init__(   self  :    Optional[int]  ,     snake_case__  :    List[Any]=0.01  ,     snake_case__  :    List[Any]=1_0_0_0				):

             '''simple docstring'''






             lowercase		:List[str]       							=				p_stop
             lowercase		:List[Any]       							=				max_length







       def __iter__(   self  :    Any				):

             '''simple docstring'''






             lowercase		:Optional[int]       							=				0
             lowercase		:List[Any]       							=				False
             while not stop and count < self.max_length:
                   yield count
                   count += 1
                   lowercase		:str       							=				random.random() < self.p_stop







class 					__magic_name__		(			unittest.TestCase     ):






       def  __snake_case						(   self  :    int  ,     snake_case__  :    Union[str, Any]  ,     snake_case__  :    List[str]  ,     snake_case__  :    Tuple=False  ,     snake_case__  :    Any=True				):

             '''simple docstring'''






             lowercase		:int       							=				[
                 BatchSamplerShard(snake_case__  ,     2  ,     snake_case__  ,     split_batches=snake_case__  ,     even_batches=snake_case__				)
                 for i in range(2				)
             ]
             lowercase		:List[str]       							=				[list(snake_case__				) for batch_sampler_shard in batch_sampler_shards]
             if not split_batches:
                   self.assertListEqual([len(snake_case__				) for shard in batch_sampler_shards]  ,     [len(snake_case__				) for e in expected]				)
             self.assertListEqual(snake_case__  ,     snake_case__				)






       def  __snake_case						(   self  :    int				):

             '''simple docstring'''






             lowercase		:Tuple       							=				BatchSampler(range(2_4				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             lowercase		:Optional[int]       							=				[
                 [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
                 [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1, 2_2, 2_3]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__				)

             lowercase		:int       							=				BatchSampler(range(2_4				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             # Expected shouldn't change
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__				)

             # Check the shards when the dataset is a round multiple of batch size but not total batch size.
             lowercase		:List[Any]       							=				BatchSampler(range(2_1				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             lowercase		:Optional[Any]       							=				[
                 [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
                 [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [0, 1, 2]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__				)

             lowercase		:List[Any]       							=				BatchSampler(range(2_1				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             lowercase		:Tuple       							=				[
                 [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
                 [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__				)

             # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
             # num_processes batch.
             lowercase		:Optional[int]       							=				BatchSampler(range(2_2				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             lowercase		:Union[str, Any]       							=				[
                 [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
                 [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1, 0, 1]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__				)

             lowercase		:Union[str, Any]       							=				BatchSampler(range(2_2				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             lowercase		:List[Any]       							=				[
                 [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
                 [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     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.
             lowercase		:Any       							=				BatchSampler(range(2_0				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             lowercase		:str       							=				[
                 [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 0]],
                 [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [1, 2, 3]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__				)

             lowercase		:str       							=				BatchSampler(range(2_0				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             lowercase		:Optional[int]       							=				[
                 [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
                 [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__				)

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

             lowercase		:Dict       							=				BatchSampler(range(2				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             lowercase		:Union[str, Any]       							=				[[], []]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__				)






       def  __snake_case						(   self  :    Tuple				):

             '''simple docstring'''






             lowercase		:str       							=				BatchSampler(range(2_4				)  ,     batch_size=4  ,     drop_last=snake_case__				)
             lowercase		:Dict       							=				[
                 [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
                 [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [2_2, 2_3]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     split_batches=snake_case__				)

             lowercase		:int       							=				BatchSampler(range(2_4				)  ,     batch_size=4  ,     drop_last=snake_case__				)
             # Expected shouldn't change
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     split_batches=snake_case__				)

             # Check the shards when the dataset is not a round multiple of batch size.
             lowercase		:List[Any]       							=				BatchSampler(range(2_2				)  ,     batch_size=4  ,     drop_last=snake_case__				)
             lowercase		:int       							=				[
                 [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
                 [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [0, 1]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     split_batches=snake_case__				)

             lowercase		:List[Any]       							=				BatchSampler(range(2_2				)  ,     batch_size=4  ,     drop_last=snake_case__				)
             lowercase		:Union[str, Any]       							=				[
                 [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
                 [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     split_batches=snake_case__				)

             # Check the shards when the dataset is not a round multiple of batch size or num_processes.
             lowercase		:Any       							=				BatchSampler(range(2_1				)  ,     batch_size=4  ,     drop_last=snake_case__				)
             lowercase		:Tuple       							=				[
                 [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 0]],
                 [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [1, 2]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     split_batches=snake_case__				)

             lowercase		:Tuple       							=				BatchSampler(range(2_1				)  ,     batch_size=4  ,     drop_last=snake_case__				)
             lowercase		:Tuple       							=				[
                 [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
                 [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     split_batches=snake_case__				)

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

             lowercase		:List[Any]       							=				BatchSampler(range(2				)  ,     batch_size=4  ,     drop_last=snake_case__				)
             lowercase		:Optional[Any]       							=				[[], []]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     split_batches=snake_case__				)






       def  __snake_case						(   self  :    str				):

             '''simple docstring'''






             lowercase		:Optional[int]       							=				BatchSampler(range(2_4				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             lowercase		:Union[str, Any]       							=				[
                 [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
                 [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1, 2_2, 2_3]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     even_batches=snake_case__				)

             lowercase		:Optional[Any]       							=				BatchSampler(range(2_4				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             # Expected shouldn't change
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     even_batches=snake_case__				)

             # Check the shards when the dataset is a round multiple of batch size but not total batch size.
             lowercase		:str       							=				BatchSampler(range(2_1				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             lowercase		:Union[str, Any]       							=				[
                 [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
                 [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     even_batches=snake_case__				)

             lowercase		:Optional[int]       							=				BatchSampler(range(2_1				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             lowercase		:str       							=				[
                 [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
                 [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     even_batches=snake_case__				)

             # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
             # num_processes batch.
             lowercase		:Optional[Any]       							=				BatchSampler(range(2_2				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             lowercase		:str       							=				[
                 [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
                 [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     even_batches=snake_case__				)

             lowercase		:Any       							=				BatchSampler(range(2_2				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             lowercase		:Tuple       							=				[
                 [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
                 [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     even_batches=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.
             lowercase		:List[str]       							=				BatchSampler(range(2_0				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             lowercase		:List[Any]       							=				[
                 [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9]],
                 [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     even_batches=snake_case__				)

             lowercase		:Optional[int]       							=				BatchSampler(range(2_0				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             lowercase		:Optional[int]       							=				[
                 [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
                 [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     even_batches=snake_case__				)

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

             lowercase		:int       							=				BatchSampler(range(2				)  ,     batch_size=3  ,     drop_last=snake_case__				)
             lowercase		:Optional[Any]       							=				[[], []]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     even_batches=snake_case__				)






       def  __snake_case						(   self  :    Optional[int]				):

             '''simple docstring'''






             lowercase		:int       							=				BatchSampler(range(2_4				)  ,     batch_size=4  ,     drop_last=snake_case__				)
             lowercase		:Union[str, Any]       							=				[
                 [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
                 [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [2_2, 2_3]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     split_batches=snake_case__  ,     even_batches=snake_case__				)

             lowercase		:int       							=				BatchSampler(range(2_4				)  ,     batch_size=4  ,     drop_last=snake_case__				)
             # Expected shouldn't change
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     split_batches=snake_case__  ,     even_batches=snake_case__				)

             # Check the shards when the dataset is not a round multiple of batch size.
             lowercase		:Union[str, Any]       							=				BatchSampler(range(2_2				)  ,     batch_size=4  ,     drop_last=snake_case__				)
             lowercase		:Optional[int]       							=				[
                 [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
                 [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     split_batches=snake_case__  ,     even_batches=snake_case__				)

             lowercase		:Dict       							=				BatchSampler(range(2_2				)  ,     batch_size=4  ,     drop_last=snake_case__				)
             lowercase		:str       							=				[
                 [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
                 [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     split_batches=snake_case__  ,     even_batches=snake_case__				)

             # Check the shards when the dataset is not a round multiple of batch size or num_processes.
             lowercase		:int       							=				BatchSampler(range(2_1				)  ,     batch_size=4  ,     drop_last=snake_case__				)
             lowercase		:Any       							=				[
                 [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0]],
                 [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     split_batches=snake_case__  ,     even_batches=snake_case__				)

             lowercase		:int       							=				BatchSampler(range(2_1				)  ,     batch_size=4  ,     drop_last=snake_case__				)
             lowercase		:Tuple       							=				[
                 [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
                 [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
             ]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     split_batches=snake_case__  ,     even_batches=snake_case__				)

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

             lowercase		:Optional[int]       							=				BatchSampler(range(2				)  ,     batch_size=4  ,     drop_last=snake_case__				)
             lowercase		:Optional[Any]       							=				[[], []]
             self.check_batch_sampler_shards(snake_case__  ,     snake_case__  ,     split_batches=snake_case__  ,     even_batches=snake_case__				)






       def  __snake_case						(   self  :    Tuple				):

             '''simple docstring'''






             lowercase		:Any       							=				[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 1_0, 1_1], [1_2, 1_3]]
             lowercase		:Any       							=				[BatchSamplerShard(snake_case__  ,     2  ,     snake_case__  ,     even_batches=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], [1_2, 1_3]]				)
             self.assertListEqual(list(batch_sampler_shards[1]				)  ,     [[3, 4], [9, 1_0, 1_1]]				)






       def  __snake_case						(   self  :    Optional[int]  ,     snake_case__  :    Optional[int]  ,     snake_case__  :    int  ,     snake_case__  :    Optional[int]  ,     snake_case__  :    List[Any]=False  ,     snake_case__  :    Tuple=2  ,     snake_case__  :    List[Any]=False				):

             '''simple docstring'''






             random.seed(snake_case__				)
             lowercase		:Dict       							=				list(snake_case__				)

             lowercase		:Any       							=				[
                 IterableDatasetShard(
                     snake_case__  ,     batch_size=snake_case__  ,     drop_last=snake_case__  ,     num_processes=snake_case__  ,     process_index=snake_case__  ,     split_batches=snake_case__  ,     )
                 for i in range(snake_case__				)
             ]
             lowercase		:Optional[Any]       							=				[]
             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(snake_case__				)
                   iterable_dataset_lists.append(list(snake_case__				)				)

             lowercase		:Optional[Any]       							=				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
             lowercase		:Union[str, Any]       							=				iterable_dataset_lists[0]
             for l in iterable_dataset_lists[1:]:
                   self.assertEqual(len(snake_case__				)  ,     len(snake_case__				)				)
                   self.assertTrue(len(snake_case__				) % shard_batch_size == 0				)

             lowercase		:List[str]       							=				[]
             for idx in range(0  ,     len(snake_case__				)  ,     snake_case__				):
                   for l in iterable_dataset_lists:
                         observed += l[idx : idx + shard_batch_size]

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






       def  __snake_case						(   self  :    Dict				):

             '''simple docstring'''






             lowercase		:Tuple       							=				4_2
             lowercase		:List[Any]       							=				RandomIterableDataset()

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

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

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






       def  __snake_case						(   self  :    List[str]				):

             '''simple docstring'''






             lowercase		:Dict       							=				BatchSampler(range(1_6				)  ,     batch_size=4  ,     drop_last=snake_case__				)
             lowercase		:str       							=				SkipBatchSampler(snake_case__  ,     2				)
             self.assertListEqual(list(snake_case__				)  ,     [[8, 9, 1_0, 1_1], [1_2, 1_3, 1_4, 1_5]]				)






       def  __snake_case						(   self  :    Optional[int]				):

             '''simple docstring'''






             lowercase		:List[str]       							=				SkipDataLoader(list(range(1_6				)				)  ,     batch_size=4  ,     skip_batches=2				)
             self.assertListEqual([t.tolist() for t in dataloader]  ,     [[8, 9, 1_0, 1_1], [1_2, 1_3, 1_4, 1_5]]				)






       def  __snake_case						(   self  :    str				):

             '''simple docstring'''






             lowercase		:Optional[Any]       							=				DataLoader(list(range(1_6				)				)  ,     batch_size=4				)
             lowercase		:Any       							=				skip_first_batches(snake_case__  ,     num_batches=2				)
             self.assertListEqual([t.tolist() for t in new_dataloader]  ,     [[8, 9, 1_0, 1_1], [1_2, 1_3, 1_4, 1_5]]				)






       def  __snake_case						(   self  :    Any				):

             '''simple docstring'''






             lowercase		:List[Any]       							=				DataLoaderShard(list(range(1_6				)				)  ,     batch_size=4				)
             for idx, _ in enumerate(snake_case__				):
                   self.assertEqual(dataloader.end_of_dataloader  ,     idx == 3				)

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







       def  __snake_case						(   self  :    Any				):

             '''simple docstring'''






             Accelerator()
             lowercase		:int       							=				DataLoaderDispatcher(range(1_6				)  ,     batch_size=4				)
             for idx, _ in enumerate(snake_case__				):
                   self.assertEqual(dataloader.end_of_dataloader  ,     idx == 3				)

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





"""simple docstring"""



import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()
class 					__magic_name__		(			__UpperCAmelCase			,						unittest.TestCase     ):
       __A			:      Tuple        =   PriorTransformer
       __A			:      Any        =   "hidden_states"






       @property
       def  __snake_case						(   self  :    str				):

              '''simple docstring'''






              lowercase		:Union[str, Any]       							=				4
              lowercase		:List[Any]       							=				8
              lowercase		:Any       							=				7

              lowercase		:Optional[Any]       							=				floats_tensor((batch_size, embedding_dim)				).to(snake_case__				)

              lowercase		:List[Any]       							=				floats_tensor((batch_size, embedding_dim)				).to(snake_case__				)
              lowercase		:Union[str, Any]       							=				floats_tensor((batch_size, num_embeddings, embedding_dim)				).to(snake_case__				)

              return {
                  "hidden_states": hidden_states,
                  "timestep": 2,
                  "proj_embedding": proj_embedding,
                  "encoder_hidden_states": encoder_hidden_states,
              }






       def  __snake_case						(   self  :    Dict  ,     snake_case__  :    Optional[Any]=0				):

              '''simple docstring'''






              torch.manual_seed(snake_case__				)
              lowercase		:Union[str, Any]       							=				4
              lowercase		:Optional[int]       							=				8
              lowercase		:str       							=				7

              lowercase		:Dict       							=				torch.randn((batch_size, embedding_dim)				).to(snake_case__				)

              lowercase		:Union[str, Any]       							=				torch.randn((batch_size, embedding_dim)				).to(snake_case__				)
              lowercase		:Tuple       							=				torch.randn((batch_size, num_embeddings, embedding_dim)				).to(snake_case__				)

              return {
                  "hidden_states": hidden_states,
                  "timestep": 2,
                  "proj_embedding": proj_embedding,
                  "encoder_hidden_states": encoder_hidden_states,
              }






       @property
       def  __snake_case						(   self  :    Union[str, Any]				):

              '''simple docstring'''






              return (4, 8)






       @property
       def  __snake_case						(   self  :    Dict				):

              '''simple docstring'''






              return (4, 8)






       def  __snake_case						(   self  :    Optional[Any]				):

              '''simple docstring'''






              lowercase		:Union[str, Any]       							=				{
                  '''num_attention_heads''': 2,
                  '''attention_head_dim''': 4,
                  '''num_layers''': 2,
                  '''embedding_dim''': 8,
                  '''num_embeddings''': 7,
                  '''additional_embeddings''': 4,
              }
              lowercase		:List[str]       							=				self.dummy_input
              return init_dict, inputs_dict






       def  __snake_case						(   self  :    List[str]				):

              '''simple docstring'''






              lowercase							,  lowercase		:Optional[int]       							=				PriorTransformer.from_pretrained(
                  '''hf-internal-testing/prior-dummy'''  ,     output_loading_info=snake_case__				)
              self.assertIsNotNone(snake_case__				)
              self.assertEqual(len(loading_info['''missing_keys''']				)  ,     0				)

              model.to(snake_case__				)
              lowercase		:List[Any]       							=				model(**self.dummy_input				)[0]

              assert hidden_states is not None, "Make sure output is not None"






       def  __snake_case						(   self  :    Union[str, Any]				):

              '''simple docstring'''






              lowercase							,  lowercase		:List[Any]       							=				self.prepare_init_args_and_inputs_for_common()

              lowercase		:str       							=				self.model_class(**snake_case__				)
              lowercase		:Optional[Any]       							=				inspect.signature(model.forward				)
              # signature.parameters is an OrderedDict => so arg_names order is deterministic
              lowercase		:List[Any]       							=				[*signature.parameters.keys()]

              lowercase		:Optional[Any]       							=				['''hidden_states''', '''timestep''']
              self.assertListEqual(arg_names[:2]  ,     snake_case__				)







       def  __snake_case						(   self  :    Optional[int]				):

              '''simple docstring'''






              lowercase		:Optional[int]       							=				PriorTransformer.from_pretrained('''hf-internal-testing/prior-dummy'''				)
              lowercase		:List[Any]       							=				model.to(snake_case__				)

              if hasattr(snake_case__  ,     '''set_default_attn_processor'''				):
                     model.set_default_attn_processor()

              lowercase		:List[str]       							=				self.get_dummy_seed_input()

              with torch.no_grad():
                     lowercase		:List[Any]       							=				model(**snake_case__				)[0]

              lowercase		:Dict       							=				output[0, :5].flatten().cpu()
              print(snake_case__				)

              # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
              # the expected output slices are not the same for CPU and GPU.
              lowercase		:str       							=				torch.tensor([-1.34_36, -0.28_70, 0.75_38, 0.43_68, -0.02_39]				)
              self.assertTrue(torch_all_close(snake_case__  ,     snake_case__  ,     rtol=1e-2				)				)







@slow
class 					__magic_name__		(			unittest.TestCase     ):






       def  __snake_case						(   self  :    Tuple  ,     snake_case__  :    int=1  ,     snake_case__  :    int=7_6_8  ,     snake_case__  :    Optional[Any]=7_7  ,     snake_case__  :    Union[str, Any]=0				):

              '''simple docstring'''






              torch.manual_seed(snake_case__				)
              lowercase		:Any       							=				batch_size
              lowercase		:Dict       							=				embedding_dim
              lowercase		:int       							=				num_embeddings

              lowercase		:int       							=				torch.randn((batch_size, embedding_dim)				).to(snake_case__				)

              lowercase		:List[str]       							=				torch.randn((batch_size, embedding_dim)				).to(snake_case__				)
              lowercase		:Optional[Any]       							=				torch.randn((batch_size, num_embeddings, embedding_dim)				).to(snake_case__				)

              return {
                  "hidden_states": hidden_states,
                  "timestep": 2,
                  "proj_embedding": proj_embedding,
                  "encoder_hidden_states": encoder_hidden_states,
              }






       def  __snake_case						(   self  :    Dict				):

              '''simple docstring'''






              super().tearDown()
              gc.collect()
              torch.cuda.empty_cache()







       @parameterized.expand(
           [
               # fmt: off
               [1_3, [-0.58_61,  0.12_83, -0.09_31,  0.08_82,  0.44_76,  0.13_29, -0.04_98,  0.06_40]],
               [3_7, [-0.49_13,  0.01_10, -0.04_83,  0.05_41,  0.49_54, -0.01_70,  0.03_54,  0.16_51]],
               # fmt: on
           ]				)
       def  __snake_case						(   self  :    Dict  ,     snake_case__  :    Dict  ,     snake_case__  :    Union[str, Any]				):

              '''simple docstring'''






              lowercase		:str       							=				PriorTransformer.from_pretrained('''kandinsky-community/kandinsky-2-1-prior'''  ,     subfolder='''prior'''				)
              model.to(snake_case__				)
              lowercase		:Union[str, Any]       							=				self.get_dummy_seed_input(seed=snake_case__				)

              with torch.no_grad():
                     lowercase		:int       							=				model(**snake_case__				)[0]

              assert list(sample.shape				) == [1, 7_6_8]

              lowercase		:List[Any]       							=				sample[0, :8].flatten().cpu()
              print(snake_case__				)
              lowercase		:Union[str, Any]       							=				torch.tensor(snake_case__				)

              assert torch_all_close(snake_case__  ,     snake_case__  ,     atol=1e-3				)

