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




		"""simple docstring"""





		def __init__(     self							,  lowerCAmelCase__=0.01							,  lowerCAmelCase__=1_0_0_0):
				__SCREAMING_SNAKE_CASE            = p_stop
				__SCREAMING_SNAKE_CASE            = max_length




		def __iter__(     self):
				__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 						SCREAMING_SNAKE_CASE_ (						unittest.TestCase ):




		"""simple docstring"""





		def 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__=False							,  lowerCAmelCase__=True):
				__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 	snake_case_     (     self):
				# 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 	snake_case_     (     self):
				# 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 	snake_case_     (     self):
				# 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 	snake_case_     (     self):
				# 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 	snake_case_     (     self):
				__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 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__=False							,  lowerCAmelCase__=2							,  lowerCAmelCase__=False):
				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 	snake_case_     (     self):
				__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 	snake_case_     (     self):
				__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 	snake_case_     (     self):
				__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 	snake_case_     (     self):
				__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 	snake_case_     (     self):
				__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 	snake_case_     (     self):
				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"""




from ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowerCamelCase							 =     logging.get_logger(__name__)

__lowerCamelCase							 =     {
    "transfo-xl-wt103": "https://huggingface.co/transfo-xl-wt103/resolve/main/config.json",
}
class 			UpperCamelCase__(	__A       ):
						lowerCAmelCase__     :  Union[str, Any]  							=						'transfo-xl'
						lowerCAmelCase__     :  Any  							=						['mems']
						lowerCAmelCase__     :  Tuple  							=						{
						    'n_token': 'vocab_size',
						    'hidden_size': 'd_model',
						    'num_attention_heads': 'n_head',
						    'num_hidden_layers': 'n_layer',
						}





						def __init__(				self						,__UpperCAmelCase=26_77_35						,__UpperCAmelCase=[2_00_00, 4_00_00, 20_00_00]						,__UpperCAmelCase=10_24						,__UpperCAmelCase=10_24						,__UpperCAmelCase=16						,__UpperCAmelCase=64						,__UpperCAmelCase=40_96						,__UpperCAmelCase=4						,__UpperCAmelCase=False						,__UpperCAmelCase=18						,__UpperCAmelCase=16_00						,__UpperCAmelCase=10_00						,__UpperCAmelCase=True						,__UpperCAmelCase=True						,__UpperCAmelCase=0						,__UpperCAmelCase=-1						,__UpperCAmelCase=True						,__UpperCAmelCase=0.1						,__UpperCAmelCase=0.0						,__UpperCAmelCase=True						,__UpperCAmelCase="normal"						,__UpperCAmelCase=0.0_1						,__UpperCAmelCase=0.0_1						,__UpperCAmelCase=0.0_2						,__UpperCAmelCase=1e-5						,__UpperCAmelCase=0						,**__UpperCAmelCase						,)    -> Tuple:
													A__		   =							vocab_size
													A__		   =							[]
													self.cutoffs.extend(__UpperCAmelCase   )
													if proj_share_all_but_first:
																				A__		   =							[False] + [True] * len(self.cutoffs   )
													else:
																				A__		   =							[False] + [False] * len(self.cutoffs   )
													A__		   =							d_model
													A__		   =							d_embed
													A__		   =							d_head
													A__		   =							d_inner
													A__		   =							div_val
													A__		   =							pre_lnorm
													A__		   =							n_layer
													A__		   =							n_head
													A__		   =							mem_len
													A__		   =							same_length
													A__		   =							attn_type
													A__		   =							clamp_len
													A__		   =							sample_softmax
													A__		   =							adaptive
													A__		   =							dropout
													A__		   =							dropatt
													A__		   =							untie_r
													A__		   =							init
													A__		   =							init_range
													A__		   =							proj_init_std
													A__		   =							init_std
													A__		   =							layer_norm_epsilon
													super().__init__(eos_token_id=__UpperCAmelCase						,**__UpperCAmelCase   )





						@property
						def  snake_case__		(				self   )    -> Optional[Any]:
													# Message copied from Transformer-XL documentation
													logger.info(f'''The model {self.model_type} is one of the few models that has no sequence length limit.'''   )
													return -1





						@max_position_embeddings.setter
						def  snake_case__		(				self						,__UpperCAmelCase   )    -> int:
													# Message copied from Transformer-XL documentation
													raise NotImplementedError(
													    f'''The model {self.model_type} is one of the few models that has no sequence length limit.'''   )





