


"""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 				UpperCamelCase_			(				UpperCamelCase):

	"""simple docstring"""





	def __init__(				self    :       Dict    ,    UpperCAmelCase__    :       List[str]=0.01    ,    UpperCAmelCase__    :       Optional[int]=1_0_0_0     ) ->			Tuple:
						__SCREAMING_SNAKE_CASE		     =       p_stop
						__SCREAMING_SNAKE_CASE		     =       max_length




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




class 				UpperCamelCase_			(				unittest.TestCase):

	"""simple docstring"""





	def 						UpperCAmelCase_   (				self    :       List[str]    ,    UpperCAmelCase__    :       Optional[Any]    ,    UpperCAmelCase__    :       Union[str, Any]    ,    UpperCAmelCase__    :       List[Any]=False    ,    UpperCAmelCase__    :       List[Any]=True     ) ->			Optional[Any]:
						__SCREAMING_SNAKE_CASE		     =       [
						    BatchSamplerShard(UpperCAmelCase__    ,    2    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )
						    for i in range(2     )
						]
						__SCREAMING_SNAKE_CASE		     =       [list(UpperCAmelCase__     ) for batch_sampler_shard in batch_sampler_shards]
						if not split_batches:
											self.assertListEqual([len(UpperCAmelCase__     ) for shard in batch_sampler_shards]    ,    [len(UpperCAmelCase__     ) for e in expected]     )
						self.assertListEqual(UpperCAmelCase__    ,    UpperCAmelCase__     )




	def 						UpperCAmelCase_   (				self    :       List[Any]     ) ->			Optional[int]:
						# Check the shards when the dataset is a round multiple of total batch size.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_4     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_4     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						# Expected shouldn't change
						self.check_batch_sampler_shards(UpperCAmelCase__    ,    UpperCAmelCase__     )

						# Check the shards when the dataset is a round multiple of batch size but not total batch size.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_1     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_1     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__     )

						# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
						# num_processes batch.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_2     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_2     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__     )

						# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
						# num_processes batch.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_0     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_0     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__     )

						# Check the shards when the dataset is very small.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [[[0, 1, 0]], [[1, 0, 1]]]
						self.check_batch_sampler_shards(UpperCAmelCase__    ,    UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [[], []]
						self.check_batch_sampler_shards(UpperCAmelCase__    ,    UpperCAmelCase__     )




	def 						UpperCAmelCase_   (				self    :       Optional[int]     ) ->			Optional[int]:
						# Check the shards when the dataset is a round multiple of batch size.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_4     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_4     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						# Expected shouldn't change
						self.check_batch_sampler_shards(UpperCAmelCase__    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__     )

						# Check the shards when the dataset is not a round multiple of batch size.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_2     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_2     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__     )

						# Check the shards when the dataset is not a round multiple of batch size or num_processes.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_1     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_1     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__     )

						# Check the shards when the dataset is very small.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [[[0, 1]], [[0, 1]]]
						self.check_batch_sampler_shards(UpperCAmelCase__    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [[], []]
						self.check_batch_sampler_shards(UpperCAmelCase__    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__     )




	def 						UpperCAmelCase_   (				self    :       List[str]     ) ->			Optional[int]:
						# Check the shards when the dataset is a round multiple of total batch size.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_4     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_4     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						# Expected shouldn't change
						self.check_batch_sampler_shards(UpperCAmelCase__    ,    UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )

						# Check the shards when the dataset is a round multiple of batch size but not total batch size.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_1     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_1     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )

						# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
						# num_processes batch.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_2     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_2     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )

						# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
						# num_processes batch.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_0     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_0     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )

						# Check the shards when the dataset is very small.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [[[0, 1]], []]
						self.check_batch_sampler_shards(UpperCAmelCase__    ,    UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2     )    ,    batch_size=3    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [[], []]
						self.check_batch_sampler_shards(UpperCAmelCase__    ,    UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )




	def 						UpperCAmelCase_   (				self    :       Optional[int]     ) ->			List[Any]:
						# Check the shards when the dataset is a round multiple of batch size.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_4     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_4     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						# Expected shouldn't change
						self.check_batch_sampler_shards(UpperCAmelCase__    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )

						# Check the shards when the dataset is not a round multiple of batch size.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_2     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_2     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )

						# Check the shards when the dataset is not a round multiple of batch size or num_processes.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_1     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2_1     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [
						    [[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(UpperCAmelCase__    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )

						# Check the shards when the dataset is very small.
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [[[0, 1]], []]
						self.check_batch_sampler_shards(UpperCAmelCase__    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )

						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(2     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       [[], []]
						self.check_batch_sampler_shards(UpperCAmelCase__    ,    UpperCAmelCase__    ,    split_batches=UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     )




	def 						UpperCAmelCase_   (				self    :       Any     ) ->			Any:
						__SCREAMING_SNAKE_CASE		     =       [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 1_0, 1_1], [1_2, 1_3]]
						__SCREAMING_SNAKE_CASE		     =       [BatchSamplerShard(UpperCAmelCase__    ,    2    ,    UpperCAmelCase__    ,    even_batches=UpperCAmelCase__     ) 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 						UpperCAmelCase_   (				self    :       Dict    ,    UpperCAmelCase__    :       Optional[int]    ,    UpperCAmelCase__    :       str    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       Optional[Any]=False    ,    UpperCAmelCase__    :       List[Any]=2    ,    UpperCAmelCase__    :       Union[str, Any]=False     ) ->			str:
						random.seed(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       list(UpperCAmelCase__     )

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

						__SCREAMING_SNAKE_CASE		     =       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
						__SCREAMING_SNAKE_CASE		     =       iterable_dataset_lists[0]
						for l in iterable_dataset_lists[1:]:
											self.assertEqual(len(UpperCAmelCase__     )    ,    len(UpperCAmelCase__     )     )
											self.assertTrue(len(UpperCAmelCase__     ) % shard_batch_size == 0     )

						__SCREAMING_SNAKE_CASE		     =       []
						for idx in range(0    ,    len(UpperCAmelCase__     )    ,    UpperCAmelCase__     ):
											for l in iterable_dataset_lists:
																observed += l[idx : idx + shard_batch_size]

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




	def 						UpperCAmelCase_   (				self    :       str     ) ->			str:
						__SCREAMING_SNAKE_CASE		     =       4_2
						__SCREAMING_SNAKE_CASE		     =       RandomIterableDataset()

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

						# Edge case with a very small dataset
						__SCREAMING_SNAKE_CASE		     =       RandomIterableDataset(max_length=2     )

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




	def 						UpperCAmelCase_   (				self    :       Tuple     ) ->			Union[str, Any]:
						__SCREAMING_SNAKE_CASE		     =       BatchSampler(range(1_6     )    ,    batch_size=4    ,    drop_last=UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       SkipBatchSampler(UpperCAmelCase__    ,    2     )
						self.assertListEqual(list(UpperCAmelCase__     )    ,    [[8, 9, 1_0, 1_1], [1_2, 1_3, 1_4, 1_5]]     )




	def 						UpperCAmelCase_   (				self    :       Dict     ) ->			List[Any]:
						__SCREAMING_SNAKE_CASE		     =       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 						UpperCAmelCase_   (				self    :       List[Any]     ) ->			Any:
						__SCREAMING_SNAKE_CASE		     =       DataLoader(list(range(1_6     )     )    ,    batch_size=4     )
						__SCREAMING_SNAKE_CASE		     =       skip_first_batches(UpperCAmelCase__    ,    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 						UpperCAmelCase_   (				self    :       Optional[Any]     ) ->			str:
						__SCREAMING_SNAKE_CASE		     =       DataLoaderShard(list(range(1_6     )     )    ,    batch_size=4     )
						for idx, _ in enumerate(UpperCAmelCase__     ):
											self.assertEqual(dataloader.end_of_dataloader    ,    idx == 3     )

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




	def 						UpperCAmelCase_   (				self    :       int     ) ->			Optional[Any]:
						Accelerator()
						__SCREAMING_SNAKE_CASE		     =       DataLoaderDispatcher(range(1_6     )    ,    batch_size=4     )
						for idx, _ in enumerate(UpperCAmelCase__     ):
											self.assertEqual(dataloader.end_of_dataloader    ,    idx == 3     )

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




"""simple docstring"""







import math
from typing import Dict, Iterable, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    is_torch_available,
    is_torch_tensor,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_torch_available():
		import torch

if is_vision_available():
		import PIL


a__       :							Dict        =  logging.get_logger(__name__)




def      UpperCAmelCase__ (lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_     ):




					'''simple docstring'''


					def constraint_to_multiple_of(lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_=0						,		lowerCAmelCase_=None     ):
										__SCREAMING_SNAKE_CASE		     =       round(val / multiple     ) * multiple

										if max_val is not None and x > max_val:
															__SCREAMING_SNAKE_CASE		     =       math.floor(val / multiple     ) * multiple

										if x < min_val:
															__SCREAMING_SNAKE_CASE		     =       math.ceil(val / multiple     ) * multiple

										return x

					__SCREAMING_SNAKE_CASE		     =       (output_size, output_size) if isinstance(lowerCAmelCase_						,		lowerCAmelCase_     ) else output_size

					__SCREAMING_SNAKE_CASE     ,						__SCREAMING_SNAKE_CASE		     =       get_image_size(lowerCAmelCase_     )
					__SCREAMING_SNAKE_CASE     ,						__SCREAMING_SNAKE_CASE		     =       output_size

					# determine new height and width
					__SCREAMING_SNAKE_CASE		     =       output_height / input_height
					__SCREAMING_SNAKE_CASE		     =       output_width / input_width

					if keep_aspect_ratio:
										# scale as little as possible
										if abs(1 - scale_width     ) < abs(1 - scale_height     ):
															# fit width
															__SCREAMING_SNAKE_CASE		     =       scale_width
										else:
															# fit height
															__SCREAMING_SNAKE_CASE		     =       scale_height

					__SCREAMING_SNAKE_CASE		     =       constraint_to_multiple_of(scale_height * input_height						,		multiple=lowerCAmelCase_     )
					__SCREAMING_SNAKE_CASE		     =       constraint_to_multiple_of(scale_width * input_width						,		multiple=lowerCAmelCase_     )

					return (new_height, new_width)




class 				UpperCamelCase_			(				UpperCamelCase):

	"""simple docstring"""





	snake_case__       :						List[Any]					      =    ["pixel_values"]




	def __init__(				self    :       int    ,    UpperCAmelCase__    :       bool = True    ,    UpperCAmelCase__    :       Dict[str, int] = None    ,    UpperCAmelCase__    :       PILImageResampling = PILImageResampling.BILINEAR    ,    UpperCAmelCase__    :       bool = False    ,    UpperCAmelCase__    :       int = 1    ,    UpperCAmelCase__    :       bool = True    ,    UpperCAmelCase__    :       Union[int, float] = 1 / 2_5_5    ,    UpperCAmelCase__    :       bool = True    ,    UpperCAmelCase__    :       Optional[Union[float, List[float]]] = None    ,    UpperCAmelCase__    :       Optional[Union[float, List[float]]] = None    ,    **UpperCAmelCase__    :       Optional[Any]    ,    ) ->			None:
						super().__init__(**UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       size if size is not None else {"height": 3_8_4, "width": 3_8_4}
						__SCREAMING_SNAKE_CASE		     =       get_size_dict(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       do_resize
						__SCREAMING_SNAKE_CASE		     =       size
						__SCREAMING_SNAKE_CASE		     =       keep_aspect_ratio
						__SCREAMING_SNAKE_CASE		     =       ensure_multiple_of
						__SCREAMING_SNAKE_CASE		     =       resample
						__SCREAMING_SNAKE_CASE		     =       do_rescale
						__SCREAMING_SNAKE_CASE		     =       rescale_factor
						__SCREAMING_SNAKE_CASE		     =       do_normalize
						__SCREAMING_SNAKE_CASE		     =       image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
						__SCREAMING_SNAKE_CASE		     =       image_std if image_std is not None else IMAGENET_STANDARD_STD




	def 						UpperCAmelCase_   (				self    :       Any    ,    UpperCAmelCase__    :       np.ndarray    ,    UpperCAmelCase__    :       Dict[str, int]    ,    UpperCAmelCase__    :       bool = False    ,    UpperCAmelCase__    :       int = 1    ,    UpperCAmelCase__    :       PILImageResampling = PILImageResampling.BICUBIC    ,    UpperCAmelCase__    :       Optional[Union[str, ChannelDimension]] = None    ,    **UpperCAmelCase__    :       Dict    ,    ) ->			np.ndarray:
						__SCREAMING_SNAKE_CASE		     =       get_size_dict(UpperCAmelCase__     )
						if "height" not in size or "width" not in size:
											raise ValueError(F"""The size dictionary must contain the keys 'height' and 'width'. Got {size.keys()}"""     )
						__SCREAMING_SNAKE_CASE		     =       get_resize_output_image_size(
						    UpperCAmelCase__    ,    output_size=(size["height"], size["width"])    ,    keep_aspect_ratio=UpperCAmelCase__    ,    multiple=UpperCAmelCase__    ,    )
						return resize(UpperCAmelCase__    ,    size=UpperCAmelCase__    ,    resample=UpperCAmelCase__    ,    data_format=UpperCAmelCase__    ,    **UpperCAmelCase__     )




	def 						UpperCAmelCase_   (				self    :       Tuple    ,    UpperCAmelCase__    :       np.ndarray    ,    UpperCAmelCase__    :       Union[int, float]    ,    UpperCAmelCase__    :       Optional[Union[str, ChannelDimension]] = None    ,    **UpperCAmelCase__    :       Union[str, Any]    ,    ) ->			Any:
						return rescale(UpperCAmelCase__    ,    scale=UpperCAmelCase__    ,    data_format=UpperCAmelCase__    ,    **UpperCAmelCase__     )




	def 						UpperCAmelCase_   (				self    :       str    ,    UpperCAmelCase__    :       np.ndarray    ,    UpperCAmelCase__    :       Union[float, List[float]]    ,    UpperCAmelCase__    :       Union[float, List[float]]    ,    UpperCAmelCase__    :       Optional[Union[str, ChannelDimension]] = None    ,    **UpperCAmelCase__    :       Optional[Any]    ,    ) ->			np.ndarray:
						return normalize(UpperCAmelCase__    ,    mean=UpperCAmelCase__    ,    std=UpperCAmelCase__    ,    data_format=UpperCAmelCase__    ,    **UpperCAmelCase__     )




	def 						UpperCAmelCase_   (				self    :       List[Any]    ,    UpperCAmelCase__    :       ImageInput    ,    UpperCAmelCase__    :       bool = None    ,    UpperCAmelCase__    :       int = None    ,    UpperCAmelCase__    :       bool = None    ,    UpperCAmelCase__    :       int = None    ,    UpperCAmelCase__    :       PILImageResampling = None    ,    UpperCAmelCase__    :       bool = None    ,    UpperCAmelCase__    :       float = None    ,    UpperCAmelCase__    :       bool = None    ,    UpperCAmelCase__    :       Optional[Union[float, List[float]]] = None    ,    UpperCAmelCase__    :       Optional[Union[float, List[float]]] = None    ,    UpperCAmelCase__    :       Optional[Union[str, TensorType]] = None    ,    UpperCAmelCase__    :       ChannelDimension = ChannelDimension.FIRST    ,    **UpperCAmelCase__    :       int    ,    ) ->			PIL.Image.Image:
						__SCREAMING_SNAKE_CASE		     =       do_resize if do_resize is not None else self.do_resize
						__SCREAMING_SNAKE_CASE		     =       size if size is not None else self.size
						__SCREAMING_SNAKE_CASE		     =       get_size_dict(UpperCAmelCase__     )
						__SCREAMING_SNAKE_CASE		     =       keep_aspect_ratio if keep_aspect_ratio is not None else self.keep_aspect_ratio
						__SCREAMING_SNAKE_CASE		     =       ensure_multiple_of if ensure_multiple_of is not None else self.ensure_multiple_of
						__SCREAMING_SNAKE_CASE		     =       resample if resample is not None else self.resample
						__SCREAMING_SNAKE_CASE		     =       do_rescale if do_rescale is not None else self.do_rescale
						__SCREAMING_SNAKE_CASE		     =       rescale_factor if rescale_factor is not None else self.rescale_factor
						__SCREAMING_SNAKE_CASE		     =       do_normalize if do_normalize is not None else self.do_normalize
						__SCREAMING_SNAKE_CASE		     =       image_mean if image_mean is not None else self.image_mean
						__SCREAMING_SNAKE_CASE		     =       image_std if image_std is not None else self.image_std

						__SCREAMING_SNAKE_CASE		     =       make_list_of_images(UpperCAmelCase__     )

						if not valid_images(UpperCAmelCase__     ):
											raise ValueError(
											    "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
											    "torch.Tensor, tf.Tensor or jax.ndarray."     )

						if do_resize and size is None or resample is None:
											raise ValueError("Size and resample must be specified if do_resize is True."     )

						if do_rescale and rescale_factor is None:
											raise ValueError("Rescale factor must be specified if do_rescale is True."     )

						if do_normalize and (image_mean is None or image_std is None):
											raise ValueError("Image mean and std must be specified if do_normalize is True."     )

						# All transformations expect numpy arrays.
						__SCREAMING_SNAKE_CASE		     =       [to_numpy_array(UpperCAmelCase__     ) for image in images]

						if do_resize:
											__SCREAMING_SNAKE_CASE		     =       [self.resize(image=UpperCAmelCase__    ,    size=UpperCAmelCase__    ,    resample=UpperCAmelCase__     ) for image in images]

						if do_rescale:
											__SCREAMING_SNAKE_CASE		     =       [self.rescale(image=UpperCAmelCase__    ,    scale=UpperCAmelCase__     ) for image in images]

						if do_normalize:
											__SCREAMING_SNAKE_CASE		     =       [self.normalize(image=UpperCAmelCase__    ,    mean=UpperCAmelCase__    ,    std=UpperCAmelCase__     ) for image in images]

						__SCREAMING_SNAKE_CASE		     =       [to_channel_dimension_format(UpperCAmelCase__    ,    UpperCAmelCase__     ) for image in images]

						__SCREAMING_SNAKE_CASE		     =       {"pixel_values": images}
						return BatchFeature(data=UpperCAmelCase__    ,    tensor_type=UpperCAmelCase__     )




	def 						UpperCAmelCase_   (				self    :       Optional[int]    ,    UpperCAmelCase__    :       List[str]    ,    UpperCAmelCase__    :       List[Tuple] = None     ) ->			Optional[Any]:
						__SCREAMING_SNAKE_CASE		     =       outputs.logits

						# Resize logits and compute semantic segmentation maps
						if target_sizes is not None:
											if len(UpperCAmelCase__     ) != len(UpperCAmelCase__     ):
																raise ValueError(
																    "Make sure that you pass in as many target sizes as the batch dimension of the logits"     )

											if is_torch_tensor(UpperCAmelCase__     ):
																__SCREAMING_SNAKE_CASE		     =       target_sizes.numpy()

											__SCREAMING_SNAKE_CASE		     =       []

											for idx in range(len(UpperCAmelCase__     )     ):
																__SCREAMING_SNAKE_CASE		     =       torch.nn.functional.interpolate(
																    logits[idx].unsqueeze(dim=0     )    ,    size=target_sizes[idx]    ,    mode="bilinear"    ,    align_corners=UpperCAmelCase__     )
																__SCREAMING_SNAKE_CASE		     =       resized_logits[0].argmax(dim=0     )
																semantic_segmentation.append(UpperCAmelCase__     )
						else:
											__SCREAMING_SNAKE_CASE		     =       logits.argmax(dim=1     )
											__SCREAMING_SNAKE_CASE		     =       [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]     )]

						return semantic_segmentation
