

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 (   __snake_case		):







   '''simple docstring'''




   def __init__(       self					:			Union[str, Any]     , __lowerCamelCase					:			Union[str, Any]=0.01     , __lowerCamelCase					:			Any=1000			):
       SCREAMING_SNAKE_CASE										=   p_stop
       SCREAMING_SNAKE_CASE										=   max_length





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

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

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

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

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

       SCREAMING_SNAKE_CASE										=   BatchSampler(range(2			)     , batch_size=3     , drop_last=__lowerCamelCase			)
       SCREAMING_SNAKE_CASE										=   [[], []]
       self.check_batch_sampler_shards(__lowerCamelCase     , __lowerCamelCase			)





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

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

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

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

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





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

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

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

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

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

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





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

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

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

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

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





   def      _snake_case		(       self					:			Dict			):
       SCREAMING_SNAKE_CASE										=   [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
       SCREAMING_SNAKE_CASE										=   [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      _snake_case		(       self					:			Union[str, Any]     , __lowerCamelCase					:			Dict     , __lowerCamelCase					:			Optional[Any]     , __lowerCamelCase					:			str     , __lowerCamelCase					:			Any=False     , __lowerCamelCase					:			Tuple=2     , __lowerCamelCase					:			Optional[Any]=False			):
       random.seed(__lowerCamelCase			)
       SCREAMING_SNAKE_CASE										=   list(__lowerCamelCase			)

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

       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(__lowerCamelCase			)     , len(__lowerCamelCase			)			)
           self.assertTrue(len(__lowerCamelCase			) % shard_batch_size == 0			)

       SCREAMING_SNAKE_CASE										=   []
       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      _snake_case		(       self					:			int			):
       SCREAMING_SNAKE_CASE										=   42
       SCREAMING_SNAKE_CASE										=   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
       SCREAMING_SNAKE_CASE										=   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      _snake_case		(       self					:			List[str]			):
       SCREAMING_SNAKE_CASE										=   BatchSampler(range(16			)     , batch_size=4     , drop_last=__lowerCamelCase			)
       SCREAMING_SNAKE_CASE										=   SkipBatchSampler(__lowerCamelCase     , 2			)
       self.assertListEqual(list(__lowerCamelCase			)     , [[8, 9, 10, 11], [12, 13, 14, 15]]			)





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


def 	__a     (      A__			:  int       ):
    if not isinstance(A__			,  A__       ):
        raise ValueError("Input must be an integer"       )
    if input_num <= 0:
        raise ValueError("Input must be positive"       )
    return sum(
        divisor for divisor in range(1			,  input_num // 2 + 1       ) if input_num % divisor == 0       )


if __name__ == "__main__":
    import doctest

    doctest.testmod()