






"""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_			(__A	):





	def __init__( self   :					str      ,      lowerCAmelCase_   :					Any=0.0_1      ,      lowerCAmelCase_   :					int=1_000		)		->     List[str]:
				UpperCAmelCase_	:					Tuple     	=   p_stop
				UpperCAmelCase_	:					int     	=   max_length





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







class      UpperCamelCase_			(unittest.TestCase	):





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





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[str]		)		->     Union[str, Any]:
				# Check the shards when the dataset is a round multiple of total batch size.
				UpperCAmelCase_	:					List[Any]     	=   BatchSampler(range(24		)      ,      batch_size=3      ,      drop_last=lowerCAmelCase_		)
				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(lowerCAmelCase_      ,      lowerCAmelCase_		)

				UpperCAmelCase_	:					str     	=   BatchSampler(range(24		)      ,      batch_size=3      ,      drop_last=lowerCAmelCase_		)
				# Expected shouldn't change
				self.check_batch_sampler_shards(lowerCAmelCase_      ,      lowerCAmelCase_		)

				# Check the shards when the dataset is a round multiple of batch size but not total batch size.
				UpperCAmelCase_	:					Optional[int]     	=   BatchSampler(range(21		)      ,      batch_size=3      ,      drop_last=lowerCAmelCase_		)
				UpperCAmelCase_	:					Dict     	=   [
				    [[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(lowerCAmelCase_      ,      lowerCAmelCase_		)

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

				# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
				# num_processes batch.
				UpperCAmelCase_	:					Any     	=   BatchSampler(range(22		)      ,      batch_size=3      ,      drop_last=lowerCAmelCase_		)
				UpperCAmelCase_	:					str     	=   [
				    [[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(lowerCAmelCase_      ,      lowerCAmelCase_		)

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

				# 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_	:					List[Any]     	=   BatchSampler(range(20		)      ,      batch_size=3      ,      drop_last=lowerCAmelCase_		)
				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(lowerCAmelCase_      ,      lowerCAmelCase_		)

				UpperCAmelCase_	:					List[str]     	=   BatchSampler(range(20		)      ,      batch_size=3      ,      drop_last=lowerCAmelCase_		)
				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(lowerCAmelCase_      ,      lowerCAmelCase_		)

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

				UpperCAmelCase_	:					List[str]     	=   BatchSampler(range(2		)      ,      batch_size=3      ,      drop_last=lowerCAmelCase_		)
				UpperCAmelCase_	:					Union[str, Any]     	=   [[], []]
				self.check_batch_sampler_shards(lowerCAmelCase_      ,      lowerCAmelCase_		)





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

				UpperCAmelCase_	:					Optional[Any]     	=   BatchSampler(range(24		)      ,      batch_size=4      ,      drop_last=lowerCAmelCase_		)
				# Expected shouldn't change
				self.check_batch_sampler_shards(lowerCAmelCase_      ,      lowerCAmelCase_      ,      split_batches=lowerCAmelCase_		)

				# Check the shards when the dataset is not a round multiple of batch size.
				UpperCAmelCase_	:					Any     	=   BatchSampler(range(22		)      ,      batch_size=4      ,      drop_last=lowerCAmelCase_		)
				UpperCAmelCase_	:					Optional[int]     	=   [
				    [[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(lowerCAmelCase_      ,      lowerCAmelCase_      ,      split_batches=lowerCAmelCase_		)

				UpperCAmelCase_	:					List[str]     	=   BatchSampler(range(22		)      ,      batch_size=4      ,      drop_last=lowerCAmelCase_		)
				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(lowerCAmelCase_      ,      lowerCAmelCase_      ,      split_batches=lowerCAmelCase_		)

				# 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=lowerCAmelCase_		)
				UpperCAmelCase_	:					int     	=   [
				    [[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(lowerCAmelCase_      ,      lowerCAmelCase_      ,      split_batches=lowerCAmelCase_		)

				UpperCAmelCase_	:					Union[str, Any]     	=   BatchSampler(range(21		)      ,      batch_size=4      ,      drop_last=lowerCAmelCase_		)
				UpperCAmelCase_	:					List[str]     	=   [
				    [[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(lowerCAmelCase_      ,      lowerCAmelCase_      ,      split_batches=lowerCAmelCase_		)

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

				UpperCAmelCase_	:					Tuple     	=   BatchSampler(range(2		)      ,      batch_size=4      ,      drop_last=lowerCAmelCase_		)
				UpperCAmelCase_	:					Optional[Any]     	=   [[], []]
				self.check_batch_sampler_shards(lowerCAmelCase_      ,      lowerCAmelCase_      ,      split_batches=lowerCAmelCase_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					int		)		->     Any:
				# Check the shards when the dataset is a round multiple of total batch size.
				UpperCAmelCase_	:					Dict     	=   BatchSampler(range(24		)      ,      batch_size=3      ,      drop_last=lowerCAmelCase_		)
				UpperCAmelCase_	:					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(lowerCAmelCase_      ,      lowerCAmelCase_      ,      even_batches=lowerCAmelCase_		)

				UpperCAmelCase_	:					int     	=   BatchSampler(range(24		)      ,      batch_size=3      ,      drop_last=lowerCAmelCase_		)
				# Expected shouldn't change
				self.check_batch_sampler_shards(lowerCAmelCase_      ,      lowerCAmelCase_      ,      even_batches=lowerCAmelCase_		)

				# Check the shards when the dataset is a round multiple of batch size but not total batch size.
				UpperCAmelCase_	:					Union[str, Any]     	=   BatchSampler(range(21		)      ,      batch_size=3      ,      drop_last=lowerCAmelCase_		)
				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]],
				]
				self.check_batch_sampler_shards(lowerCAmelCase_      ,      lowerCAmelCase_      ,      even_batches=lowerCAmelCase_		)

				UpperCAmelCase_	:					Tuple     	=   BatchSampler(range(21		)      ,      batch_size=3      ,      drop_last=lowerCAmelCase_		)
				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(lowerCAmelCase_      ,      lowerCAmelCase_      ,      even_batches=lowerCAmelCase_		)

				# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
				# num_processes batch.
				UpperCAmelCase_	:					str     	=   BatchSampler(range(22		)      ,      batch_size=3      ,      drop_last=lowerCAmelCase_		)
				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(lowerCAmelCase_      ,      lowerCAmelCase_      ,      even_batches=lowerCAmelCase_		)

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

				# 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_	:					Tuple     	=   BatchSampler(range(20		)      ,      batch_size=3      ,      drop_last=lowerCAmelCase_		)
				UpperCAmelCase_	:					str     	=   [
				    [[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(lowerCAmelCase_      ,      lowerCAmelCase_      ,      even_batches=lowerCAmelCase_		)

				UpperCAmelCase_	:					str     	=   BatchSampler(range(20		)      ,      batch_size=3      ,      drop_last=lowerCAmelCase_		)
				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(lowerCAmelCase_      ,      lowerCAmelCase_      ,      even_batches=lowerCAmelCase_		)

				# Check the shards when the dataset is very small.
				UpperCAmelCase_	:					Tuple     	=   BatchSampler(range(2		)      ,      batch_size=3      ,      drop_last=lowerCAmelCase_		)
				UpperCAmelCase_	:					Optional[int]     	=   [[[0, 1]], []]
				self.check_batch_sampler_shards(lowerCAmelCase_      ,      lowerCAmelCase_      ,      even_batches=lowerCAmelCase_		)

				UpperCAmelCase_	:					str     	=   BatchSampler(range(2		)      ,      batch_size=3      ,      drop_last=lowerCAmelCase_		)
				UpperCAmelCase_	:					Any     	=   [[], []]
				self.check_batch_sampler_shards(lowerCAmelCase_      ,      lowerCAmelCase_      ,      even_batches=lowerCAmelCase_		)





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

				UpperCAmelCase_	:					Optional[int]     	=   BatchSampler(range(24		)      ,      batch_size=4      ,      drop_last=lowerCAmelCase_		)
				# Expected shouldn't change
				self.check_batch_sampler_shards(lowerCAmelCase_      ,      lowerCAmelCase_      ,      split_batches=lowerCAmelCase_      ,      even_batches=lowerCAmelCase_		)

				# Check the shards when the dataset is not a round multiple of batch size.
				UpperCAmelCase_	:					Any     	=   BatchSampler(range(22		)      ,      batch_size=4      ,      drop_last=lowerCAmelCase_		)
				UpperCAmelCase_	:					Dict     	=   [
				    [[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(lowerCAmelCase_      ,      lowerCAmelCase_      ,      split_batches=lowerCAmelCase_      ,      even_batches=lowerCAmelCase_		)

				UpperCAmelCase_	:					Optional[int]     	=   BatchSampler(range(22		)      ,      batch_size=4      ,      drop_last=lowerCAmelCase_		)
				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(lowerCAmelCase_      ,      lowerCAmelCase_      ,      split_batches=lowerCAmelCase_      ,      even_batches=lowerCAmelCase_		)

				# Check the shards when the dataset is not a round multiple of batch size or num_processes.
				UpperCAmelCase_	:					Any     	=   BatchSampler(range(21		)      ,      batch_size=4      ,      drop_last=lowerCAmelCase_		)
				UpperCAmelCase_	:					int     	=   [
				    [[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(lowerCAmelCase_      ,      lowerCAmelCase_      ,      split_batches=lowerCAmelCase_      ,      even_batches=lowerCAmelCase_		)

				UpperCAmelCase_	:					Tuple     	=   BatchSampler(range(21		)      ,      batch_size=4      ,      drop_last=lowerCAmelCase_		)
				UpperCAmelCase_	:					Union[str, 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(lowerCAmelCase_      ,      lowerCAmelCase_      ,      split_batches=lowerCAmelCase_      ,      even_batches=lowerCAmelCase_		)

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

				UpperCAmelCase_	:					Union[str, Any]     	=   BatchSampler(range(2		)      ,      batch_size=4      ,      drop_last=lowerCAmelCase_		)
				UpperCAmelCase_	:					List[str]     	=   [[], []]
				self.check_batch_sampler_shards(lowerCAmelCase_      ,      lowerCAmelCase_      ,      split_batches=lowerCAmelCase_      ,      even_batches=lowerCAmelCase_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Any		)		->     Any:
				UpperCAmelCase_	:					Optional[Any]     	=   [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
				UpperCAmelCase_	:					Optional[int]     	=   [BatchSamplerShard(lowerCAmelCase_      ,      2      ,      lowerCAmelCase_      ,      even_batches=lowerCAmelCase_		) 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 	_SCREAMING_SNAKE_CASE ( self   :					Optional[int]      ,      lowerCAmelCase_   :					Union[str, Any]      ,      lowerCAmelCase_   :					Dict      ,      lowerCAmelCase_   :					Optional[int]      ,      lowerCAmelCase_   :					Optional[int]=False      ,      lowerCAmelCase_   :					Any=2      ,      lowerCAmelCase_   :					List[Any]=False		)		->     Dict:
				random.seed(lowerCAmelCase_		)
				UpperCAmelCase_	:					List[Any]     	=   list(lowerCAmelCase_		)

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

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

				UpperCAmelCase_	:					Optional[int]     	=   []
				for idx in range(0      ,      len(lowerCAmelCase_		)      ,      lowerCAmelCase_		):
							for l in iterable_dataset_lists:
										observed += l[idx : idx + shard_batch_size]

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





	def 	_SCREAMING_SNAKE_CASE ( self   :					str		)		->     List[str]:
				UpperCAmelCase_	:					Optional[Any]     	=   42
				UpperCAmelCase_	:					Optional[int]     	=   RandomIterableDataset()

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

				# Edge case with a very small dataset
				UpperCAmelCase_	:					List[str]     	=   RandomIterableDataset(max_length=2		)

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





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[Any]		)		->     Union[str, Any]:
				UpperCAmelCase_	:					Any     	=   BatchSampler(range(16		)      ,      batch_size=4      ,      drop_last=lowerCAmelCase_		)
				UpperCAmelCase_	:					List[str]     	=   SkipBatchSampler(lowerCAmelCase_      ,      2		)
				self.assertListEqual(list(lowerCAmelCase_		)      ,      [[8, 9, 10, 11], [12, 13, 14, 15]]		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Tuple		)		->     Any:
				UpperCAmelCase_	:					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 	_SCREAMING_SNAKE_CASE ( self   :					List[str]		)		->     List[Any]:
				UpperCAmelCase_	:					List[str]     	=   DataLoader(list(range(16		)		)      ,      batch_size=4		)
				UpperCAmelCase_	:					int     	=   skip_first_batches(lowerCAmelCase_      ,      num_batches=2		)
				self.assertListEqual([t.tolist() for t in new_dataloader]      ,      [[8, 9, 10, 11], [12, 13, 14, 15]]		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[str]		)		->     str:
				UpperCAmelCase_	:					List[Any]     	=   DataLoaderShard(list(range(16		)		)      ,      batch_size=4		)
				for idx, _ in enumerate(lowerCAmelCase_		):
							self.assertEqual(dataloader.end_of_dataloader      ,      idx == 3		)

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





	def 	_SCREAMING_SNAKE_CASE ( self   :					Dict		)		->     int:
				Accelerator()
				UpperCAmelCase_	:					Optional[Any]     	=   DataLoaderDispatcher(range(16		)      ,      batch_size=4		)
				for idx, _ in enumerate(lowerCAmelCase_		):
							self.assertEqual(dataloader.end_of_dataloader      ,      idx == 3		)

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











"""simple docstring"""
import unittest

import numpy as np

from transformers import DistilBertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
			import jax.numpy as jnp

			from transformers.models.distilbert.modeling_flax_distilbert import (
			    FlaxDistilBertForMaskedLM,
			    FlaxDistilBertForMultipleChoice,
			    FlaxDistilBertForQuestionAnswering,
			    FlaxDistilBertForSequenceClassification,
			    FlaxDistilBertForTokenClassification,
			    FlaxDistilBertModel,
			)







class      UpperCamelCase_			(unittest.TestCase	):





	def __init__( self   :					int      ,      lowerCAmelCase_   :					Any      ,      lowerCAmelCase_   :					Union[str, Any]=13      ,      lowerCAmelCase_   :					Union[str, Any]=7      ,      lowerCAmelCase_   :					Optional[Any]=True      ,      lowerCAmelCase_   :					List[Any]=True      ,      lowerCAmelCase_   :					Any=True      ,      lowerCAmelCase_   :					Tuple=True      ,      lowerCAmelCase_   :					Tuple=99      ,      lowerCAmelCase_   :					str=32      ,      lowerCAmelCase_   :					Optional[int]=5      ,      lowerCAmelCase_   :					List[Any]=4      ,      lowerCAmelCase_   :					List[Any]=37      ,      lowerCAmelCase_   :					str="gelu"      ,      lowerCAmelCase_   :					Optional[Any]=0.1      ,      lowerCAmelCase_   :					Union[str, Any]=0.1      ,      lowerCAmelCase_   :					List[Any]=512      ,      lowerCAmelCase_   :					Dict=16      ,      lowerCAmelCase_   :					int=2      ,      lowerCAmelCase_   :					Union[str, Any]=0.0_2      ,      lowerCAmelCase_   :					Optional[Any]=4      ,      )		->     Tuple:
				UpperCAmelCase_	:					Optional[Any]     	=   parent
				UpperCAmelCase_	:					Tuple     	=   batch_size
				UpperCAmelCase_	:					str     	=   seq_length
				UpperCAmelCase_	:					Optional[int]     	=   is_training
				UpperCAmelCase_	:					Dict     	=   use_attention_mask
				UpperCAmelCase_	:					int     	=   use_token_type_ids
				UpperCAmelCase_	:					Tuple     	=   use_labels
				UpperCAmelCase_	:					Tuple     	=   vocab_size
				UpperCAmelCase_	:					Union[str, Any]     	=   hidden_size
				UpperCAmelCase_	:					int     	=   num_hidden_layers
				UpperCAmelCase_	:					List[str]     	=   num_attention_heads
				UpperCAmelCase_	:					str     	=   intermediate_size
				UpperCAmelCase_	:					List[str]     	=   hidden_act
				UpperCAmelCase_	:					Tuple     	=   hidden_dropout_prob
				UpperCAmelCase_	:					Optional[Any]     	=   attention_probs_dropout_prob
				UpperCAmelCase_	:					Optional[int]     	=   max_position_embeddings
				UpperCAmelCase_	:					Optional[Any]     	=   type_vocab_size
				UpperCAmelCase_	:					Dict     	=   type_sequence_label_size
				UpperCAmelCase_	:					List[Any]     	=   initializer_range
				UpperCAmelCase_	:					Any     	=   num_choices





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[str]		)		->     str:
				UpperCAmelCase_	:					List[str]     	=   ids_tensor([self.batch_size, self.seq_length]      ,      self.vocab_size		)

				UpperCAmelCase_	:					Tuple     	=   None
				if self.use_attention_mask:
							UpperCAmelCase_	:					Dict     	=   random_attention_mask([self.batch_size, self.seq_length]		)

				UpperCAmelCase_	:					Any     	=   DistilBertConfig(
				    vocab_size=self.vocab_size      ,      dim=self.hidden_size      ,      n_layers=self.num_hidden_layers      ,      n_heads=self.num_attention_heads      ,      hidden_dim=self.intermediate_size      ,      hidden_act=self.hidden_act      ,      dropout=self.hidden_dropout_prob      ,      attention_dropout=self.attention_probs_dropout_prob      ,      max_position_embeddings=self.max_position_embeddings      ,      initializer_range=self.initializer_range      ,      tie_weights_=lowerCAmelCase_      ,      )

				return config, input_ids, attention_mask





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[Any]		)		->     Union[str, Any]:
				UpperCAmelCase_	:					Any     	=   self.prepare_config_and_inputs()
				UpperCAmelCase_	,      UpperCAmelCase_	,      UpperCAmelCase_	:					List[str]     	=   config_and_inputs
				UpperCAmelCase_	:					Optional[Any]     	=   {"input_ids": input_ids, "attention_mask": attention_mask}
				return config, inputs_dict







@require_flax
class      UpperCamelCase_			(__A							,		unittest.TestCase	):
	__magic_name__       =							(
	    (
	        FlaxDistilBertModel,
	        FlaxDistilBertForMaskedLM,
	        FlaxDistilBertForMultipleChoice,
	        FlaxDistilBertForQuestionAnswering,
	        FlaxDistilBertForSequenceClassification,
	        FlaxDistilBertForTokenClassification,
	        FlaxDistilBertForQuestionAnswering,
	    )
	    if is_flax_available()
	    else ()
	)





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[str]		)		->     Optional[Any]:
				UpperCAmelCase_	:					List[Any]     	=   FlaxDistilBertModelTester(self		)





	@slow
	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[Any]		)		->     List[str]:
				for model_class_name in self.all_model_classes:
							UpperCAmelCase_	:					str     	=   model_class_name.from_pretrained("distilbert-base-uncased"		)
							UpperCAmelCase_	:					List[str]     	=   model(np.ones((1, 1)		)		)
							self.assertIsNotNone(lowerCAmelCase_		)







@require_flax
class      UpperCamelCase_			(unittest.TestCase	):





	@slow
	def 	_SCREAMING_SNAKE_CASE ( self   :					str		)		->     Optional[int]:
				UpperCAmelCase_	:					Any     	=   FlaxDistilBertModel.from_pretrained("distilbert-base-uncased"		)
				UpperCAmelCase_	:					int     	=   np.array([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]		)
				UpperCAmelCase_	:					List[str]     	=   np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]		)
				UpperCAmelCase_	:					Any     	=   model(lowerCAmelCase_      ,      attention_mask=lowerCAmelCase_		)[0]
				UpperCAmelCase_	:					Union[str, Any]     	=   (1, 11, 768)
				self.assertEqual(output.shape      ,      lowerCAmelCase_		)
				UpperCAmelCase_	:					Dict     	=   np.array([[[-0.1_6_3_9, 0.3_2_9_9, 0.1_6_4_8], [-0.1_7_4_6, 0.3_2_8_9, 0.1_7_1_0], [-0.1_8_8_4, 0.3_3_5_7, 0.1_8_1_0]]]		)

				self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]      ,      lowerCAmelCase_      ,      atol=1e-4		)		)



