


'''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 						a__							(				a__    ):





							'''simple docstring'''






							def __init__(						self   ,       lowerCamelCase_=0.01   ,       lowerCamelCase_=10_00  )			->				Union[str, Any]:
														lowerCAmelCase__			  =      p_stop
														lowerCAmelCase__			  =      max_length







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




class 						a__							(				unittest.TestCase    ):





							'''simple docstring'''






							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_=False   ,       lowerCamelCase_=True  )			->				Optional[Any]:
														lowerCAmelCase__			  =      [
														    BatchSamplerShard(lowerCamelCase_   ,       2   ,       lowerCamelCase_   ,       split_batches=lowerCamelCase_   ,       even_batches=lowerCamelCase_  )
														    for i in range(2  )
														]
														lowerCAmelCase__			  =      [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  )			->				Optional[int]:
														# Check the shards when the dataset is a round multiple of total batch size.
														lowerCAmelCase__			  =      BatchSampler(range(24  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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_  )

														lowerCAmelCase__			  =      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.
														lowerCAmelCase__			  =      BatchSampler(range(21  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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_  )

														lowerCAmelCase__			  =      BatchSampler(range(21  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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.
														lowerCAmelCase__			  =      BatchSampler(range(22  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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_  )

														lowerCAmelCase__			  =      BatchSampler(range(22  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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.
														lowerCAmelCase__			  =      BatchSampler(range(20  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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_  )

														lowerCAmelCase__			  =      BatchSampler(range(20  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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.
														lowerCAmelCase__			  =      BatchSampler(range(2  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [[[0, 1, 0]], [[1, 0, 1]]]
														self.check_batch_sampler_shards(lowerCamelCase_   ,       lowerCamelCase_  )

														lowerCAmelCase__			  =      BatchSampler(range(2  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [[], []]
														self.check_batch_sampler_shards(lowerCamelCase_   ,       lowerCamelCase_  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														# Check the shards when the dataset is a round multiple of batch size.
														lowerCAmelCase__			  =      BatchSampler(range(24  )   ,       batch_size=4   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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_  )

														lowerCAmelCase__			  =      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.
														lowerCAmelCase__			  =      BatchSampler(range(22  )   ,       batch_size=4   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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_  )

														lowerCAmelCase__			  =      BatchSampler(range(22  )   ,       batch_size=4   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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.
														lowerCAmelCase__			  =      BatchSampler(range(21  )   ,       batch_size=4   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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_  )

														lowerCAmelCase__			  =      BatchSampler(range(21  )   ,       batch_size=4   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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.
														lowerCAmelCase__			  =      BatchSampler(range(2  )   ,       batch_size=4   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [[[0, 1]], [[0, 1]]]
														self.check_batch_sampler_shards(lowerCamelCase_   ,       lowerCamelCase_   ,       split_batches=lowerCamelCase_  )

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







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														# Check the shards when the dataset is a round multiple of total batch size.
														lowerCAmelCase__			  =      BatchSampler(range(24  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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_  )

														lowerCAmelCase__			  =      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.
														lowerCAmelCase__			  =      BatchSampler(range(21  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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_  )

														lowerCAmelCase__			  =      BatchSampler(range(21  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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.
														lowerCAmelCase__			  =      BatchSampler(range(22  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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_  )

														lowerCAmelCase__			  =      BatchSampler(range(22  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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.
														lowerCAmelCase__			  =      BatchSampler(range(20  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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_  )

														lowerCAmelCase__			  =      BatchSampler(range(20  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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.
														lowerCAmelCase__			  =      BatchSampler(range(2  )   ,       batch_size=3   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [[[0, 1]], []]
														self.check_batch_sampler_shards(lowerCamelCase_   ,       lowerCamelCase_   ,       even_batches=lowerCamelCase_  )

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







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														# Check the shards when the dataset is a round multiple of batch size.
														lowerCAmelCase__			  =      BatchSampler(range(24  )   ,       batch_size=4   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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_  )

														lowerCAmelCase__			  =      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.
														lowerCAmelCase__			  =      BatchSampler(range(22  )   ,       batch_size=4   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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_  )

														lowerCAmelCase__			  =      BatchSampler(range(22  )   ,       batch_size=4   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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.
														lowerCAmelCase__			  =      BatchSampler(range(21  )   ,       batch_size=4   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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_  )

														lowerCAmelCase__			  =      BatchSampler(range(21  )   ,       batch_size=4   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [
														    [[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.
														lowerCAmelCase__			  =      BatchSampler(range(2  )   ,       batch_size=4   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      [[[0, 1]], []]
														self.check_batch_sampler_shards(lowerCamelCase_   ,       lowerCamelCase_   ,       split_batches=lowerCamelCase_   ,       even_batches=lowerCamelCase_  )

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







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[int]:
														lowerCAmelCase__			  =      [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
														lowerCAmelCase__			  =      [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   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_   ,       lowerCamelCase_=False   ,       lowerCamelCase_=2   ,       lowerCamelCase_=False  )			->				str:
														random.seed(lowerCamelCase_  )
														lowerCAmelCase__			  =      list(lowerCamelCase_  )

														lowerCAmelCase__			  =      [
														    IterableDatasetShard(
														        lowerCamelCase_   ,       batch_size=lowerCamelCase_   ,       drop_last=lowerCamelCase_   ,       num_processes=lowerCamelCase_   ,       process_index=lowerCamelCase_   ,       split_batches=lowerCamelCase_   ,       )
														    for i in range(lowerCamelCase_  )
														]
														lowerCAmelCase__			  =      []
														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_  )  )

														lowerCAmelCase__			  =      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
														lowerCAmelCase__			  =      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  )

														lowerCAmelCase__			  =      []
														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  )			->				Tuple:
														lowerCAmelCase__			  =      42
														lowerCAmelCase__			  =      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
														lowerCAmelCase__			  =      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  )			->				Any:
														lowerCAmelCase__			  =      BatchSampler(range(16  )   ,       batch_size=4   ,       drop_last=lowerCamelCase_  )
														lowerCAmelCase__			  =      SkipBatchSampler(lowerCamelCase_   ,       2  )
														self.assertListEqual(list(lowerCamelCase_  )   ,       [[8, 9, 10, 11], [12, 13, 14, 15]]  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[str]:
														lowerCAmelCase__			  =      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  )			->				Tuple:
														lowerCAmelCase__			  =      DataLoader(list(range(16  )  )   ,       batch_size=4  )
														lowerCAmelCase__			  =      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[Any]:
														lowerCAmelCase__			  =      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  )			->				str:
														Accelerator()
														lowerCAmelCase__			  =      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'''






from __future__ import annotations


def    _snake_case  (	A			)		->       bool:
							lowerCAmelCase__			  =      str(A			)
							return len(A			) == 9 and set(A			) == set('''123456789'''			)


def    _snake_case  (	)		->       int | None:
							for base_num in range(9999					,       4999					,       -1			):
														lowerCAmelCase__			  =      100002 * base_num
														if is_9_pandigital(A			):
																					return candidate

							for base_num in range(333					,       99					,       -1			):
														lowerCAmelCase__			  =      1002003 * base_num
														if is_9_pandigital(A			):
																					return candidate

							return None


if __name__ == "__main__":
		print(f"""{solution() = }""")