




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       ,_SCREAMING_SNAKE_CASE=0.01       ,_SCREAMING_SNAKE_CASE=1_000  )			->		int:
         UpperCAmelCase_      :							Union[str, Any]       =      p_stop
         UpperCAmelCase_      :							Optional[Any]       =      max_length





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


class   __a(					unittest.TestCase       ):






    """simple docstring"""





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





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

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

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

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

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

         UpperCAmelCase_      :							Dict       =      BatchSampler(range(22  )       ,batch_size=3       ,drop_last=_SCREAMING_SNAKE_CASE  )
         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(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )

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

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

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

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





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

         UpperCAmelCase_      :							int       =      BatchSampler(range(24  )       ,batch_size=4       ,drop_last=_SCREAMING_SNAKE_CASE  )
         # Expected shouldn't change
         self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,split_batches=_SCREAMING_SNAKE_CASE  )

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

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

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

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

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

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





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

         UpperCAmelCase_      :							List[str]       =      BatchSampler(range(24  )       ,batch_size=3       ,drop_last=_SCREAMING_SNAKE_CASE  )
         # Expected shouldn't change
         self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,even_batches=_SCREAMING_SNAKE_CASE  )

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

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

         # 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=_SCREAMING_SNAKE_CASE  )
         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(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,even_batches=_SCREAMING_SNAKE_CASE  )

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

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

         UpperCAmelCase_      :							Optional[Any]       =      BatchSampler(range(20  )       ,batch_size=3       ,drop_last=_SCREAMING_SNAKE_CASE  )
         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(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,even_batches=_SCREAMING_SNAKE_CASE  )

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

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





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

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

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

         UpperCAmelCase_      :							Union[str, Any]       =      BatchSampler(range(22  )       ,batch_size=4       ,drop_last=_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Dict       =      [
             [[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(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,split_batches=_SCREAMING_SNAKE_CASE       ,even_batches=_SCREAMING_SNAKE_CASE  )

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

         UpperCAmelCase_      :							Dict       =      BatchSampler(range(21  )       ,batch_size=4       ,drop_last=_SCREAMING_SNAKE_CASE  )
         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(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,split_batches=_SCREAMING_SNAKE_CASE       ,even_batches=_SCREAMING_SNAKE_CASE  )

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

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





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

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

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

         UpperCAmelCase_      :							List[str]       =      []
         for idx in range(0       ,len(_SCREAMING_SNAKE_CASE  )       ,_SCREAMING_SNAKE_CASE  ):
              for l in iterable_dataset_lists:
                   observed += l[idx : idx + shard_batch_size]

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





    def 							a__				(				self  )			->		int:
         UpperCAmelCase_      :							Union[str, Any]       =      42
         UpperCAmelCase_      :							Tuple       =      RandomIterableDataset()

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

         # Edge case with a very small dataset
         UpperCAmelCase_      :							Dict       =      RandomIterableDataset(max_length=2  )

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





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





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





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

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





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

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





from math import ceil







def lowerCamelCase__       (      _lowercase = 1001				):

     '''simple docstring'''


     UpperCAmelCase_      :							Optional[Any]       =      1

     for i in range(1			,						int(ceil(n / 2.0				)				)				):
          UpperCAmelCase_      :							List[Any]       =      2 * i + 1
          UpperCAmelCase_      :							Optional[int]       =      2 * i
          UpperCAmelCase_      :							Any       =      total + 4 * odd**2 - 6 * even

     return total


if __name__ == "__main__":
     import sys

     if len(sys.argv) == 1:
          print(solution())
     else:
          try:
               __a          =			int(sys.argv[1])
               print(solution(n))
          except ValueError:
               print('Invalid entry - please enter a number')