




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



 '''simple docstring'''


 def __init__(    self :  List[Any]						,							A_ :  Optional[Any]=0.01						,							A_ :  int=1000  )       ->     Union[str, Any]:

       """simple docstring"""






       lowerCamelCase_							=   p_stop
       lowerCamelCase_							=   max_length





 def __iter__(    self :  Any  )       ->     List[Any]:

       """simple docstring"""






       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     a__			(    self :  Union[str, Any]						,							A_ :  List[Any]						,							A_ :  Any						,							A_ :  List[Any]=False						,							A_ :  Optional[Any]=True  )       ->     Union[str, Any]:

       """simple docstring"""






       lowerCamelCase_							=   [
           BatchSamplerShard(A_						,							2						,							A_						,							split_batches=A_						,							even_batches=A_  )
           for i in range(2  )
       ]
       lowerCamelCase_							=   [list(A_  ) for batch_sampler_shard in batch_sampler_shards]
       if not split_batches:
             self.assertListEqual([len(A_  ) for shard in batch_sampler_shards]						,							[len(A_  ) for e in expected]  )
       self.assertListEqual(A_						,							A_  )

 def     a__			(    self :  Tuple  )       ->     int:

       """simple docstring"""






       lowerCamelCase_							=   BatchSampler(range(24  )						,							batch_size=3						,							drop_last=A_  )
       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(A_						,							A_  )

       lowerCamelCase_							=   BatchSampler(range(24  )						,							batch_size=3						,							drop_last=A_  )
       # Expected shouldn't change
       self.check_batch_sampler_shards(A_						,							A_  )

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

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

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

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

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

       lowerCamelCase_							=   BatchSampler(range(20  )						,							batch_size=3						,							drop_last=A_  )
       lowerCamelCase_							=   [
           [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
           [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
       ]
       self.check_batch_sampler_shards(A_						,							A_  )

       # Check the shards when the dataset is very small.
       lowerCamelCase_							=   BatchSampler(range(2  )						,							batch_size=3						,							drop_last=A_  )
       lowerCamelCase_							=   [[[0, 1, 0]], [[1, 0, 1]]]
       self.check_batch_sampler_shards(A_						,							A_  )

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

 def     a__			(    self :  Optional[Any]  )       ->     Optional[Any]:

       """simple docstring"""






       lowerCamelCase_							=   BatchSampler(range(24  )						,							batch_size=4						,							drop_last=A_  )
       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(A_						,							A_						,							split_batches=A_  )

       lowerCamelCase_							=   BatchSampler(range(24  )						,							batch_size=4						,							drop_last=A_  )
       # Expected shouldn't change
       self.check_batch_sampler_shards(A_						,							A_						,							split_batches=A_  )

       # Check the shards when the dataset is not a round multiple of batch size.
       lowerCamelCase_							=   BatchSampler(range(22  )						,							batch_size=4						,							drop_last=A_  )
       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(A_						,							A_						,							split_batches=A_  )

       lowerCamelCase_							=   BatchSampler(range(22  )						,							batch_size=4						,							drop_last=A_  )
       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(A_						,							A_						,							split_batches=A_  )

       # 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=A_  )
       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(A_						,							A_						,							split_batches=A_  )

       lowerCamelCase_							=   BatchSampler(range(21  )						,							batch_size=4						,							drop_last=A_  )
       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(A_						,							A_						,							split_batches=A_  )

       # Check the shards when the dataset is very small.
       lowerCamelCase_							=   BatchSampler(range(2  )						,							batch_size=4						,							drop_last=A_  )
       lowerCamelCase_							=   [[[0, 1]], [[0, 1]]]
       self.check_batch_sampler_shards(A_						,							A_						,							split_batches=A_  )

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

 def     a__			(    self :  Optional[int]  )       ->     Dict:

       """simple docstring"""






       lowerCamelCase_							=   BatchSampler(range(24  )						,							batch_size=3						,							drop_last=A_  )
       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(A_						,							A_						,							even_batches=A_  )

       lowerCamelCase_							=   BatchSampler(range(24  )						,							batch_size=3						,							drop_last=A_  )
       # Expected shouldn't change
       self.check_batch_sampler_shards(A_						,							A_						,							even_batches=A_  )

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

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

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

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

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

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

       # Check the shards when the dataset is very small.
       lowerCamelCase_							=   BatchSampler(range(2  )						,							batch_size=3						,							drop_last=A_  )
       lowerCamelCase_							=   [[[0, 1]], []]
       self.check_batch_sampler_shards(A_						,							A_						,							even_batches=A_  )

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

 def     a__			(    self :  List[Any]  )       ->     List[str]:

       """simple docstring"""






       lowerCamelCase_							=   BatchSampler(range(24  )						,							batch_size=4						,							drop_last=A_  )
       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(A_						,							A_						,							split_batches=A_						,							even_batches=A_  )

       lowerCamelCase_							=   BatchSampler(range(24  )						,							batch_size=4						,							drop_last=A_  )
       # Expected shouldn't change
       self.check_batch_sampler_shards(A_						,							A_						,							split_batches=A_						,							even_batches=A_  )

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

       lowerCamelCase_							=   BatchSampler(range(22  )						,							batch_size=4						,							drop_last=A_  )
       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(A_						,							A_						,							split_batches=A_						,							even_batches=A_  )

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

       lowerCamelCase_							=   BatchSampler(range(21  )						,							batch_size=4						,							drop_last=A_  )
       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(A_						,							A_						,							split_batches=A_						,							even_batches=A_  )

       # Check the shards when the dataset is very small.
       lowerCamelCase_							=   BatchSampler(range(2  )						,							batch_size=4						,							drop_last=A_  )
       lowerCamelCase_							=   [[[0, 1]], []]
       self.check_batch_sampler_shards(A_						,							A_						,							split_batches=A_						,							even_batches=A_  )

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

 def     a__			(    self :  Dict  )       ->     Dict:

       """simple docstring"""






       lowerCamelCase_							=   [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
       lowerCamelCase_							=   [BatchSamplerShard(A_						,							2						,							A_						,							even_batches=A_  ) 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 :  str						,							A_ :  str						,							A_ :  Optional[Any]						,							A_ :  Optional[Any]						,							A_ :  Tuple=False						,							A_ :  Optional[int]=2						,							A_ :  Optional[int]=False  )       ->     Union[str, Any]:

       """simple docstring"""






       random.seed(A_  )
       lowerCamelCase_							=   list(A_  )

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

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

       lowerCamelCase_							=   []
       for idx in range(0						,							len(A_  )						,							A_  ):
             for l in iterable_dataset_lists:
                   observed += l[idx : idx + shard_batch_size]

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

 def     a__			(    self :  List[str]  )       ->     int:

       """simple docstring"""






       lowerCamelCase_							=   42
       lowerCamelCase_							=   RandomIterableDataset()

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

       # Edge case with a very small dataset
       lowerCamelCase_							=   RandomIterableDataset(max_length=2  )

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

 def     a__			(    self :  List[str]  )       ->     List[Any]:

       """simple docstring"""






       lowerCamelCase_							=   BatchSampler(range(16  )						,							batch_size=4						,							drop_last=A_  )
       lowerCamelCase_							=   SkipBatchSampler(A_						,							2  )
       self.assertListEqual(list(A_  )						,							[[8, 9, 10, 11], [12, 13, 14, 15]]  )

 def     a__			(    self :  List[str]  )       ->     List[str]:

       """simple docstring"""






       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     a__			(    self :  Tuple  )       ->     Any:

       """simple docstring"""






       lowerCamelCase_							=   DataLoader(list(range(16  )  )						,							batch_size=4  )
       lowerCamelCase_							=   skip_first_batches(A_						,							num_batches=2  )
       self.assertListEqual([t.tolist() for t in new_dataloader]						,							[[8, 9, 10, 11], [12, 13, 14, 15]]  )

 def     a__			(    self :  str  )       ->     Any:

       """simple docstring"""






       lowerCamelCase_							=   DataLoaderShard(list(range(16  )  )						,							batch_size=4  )
       for idx, _ in enumerate(A_  ):
             self.assertEqual(dataloader.end_of_dataloader						,							idx == 3  )

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





 def     a__			(    self :  Union[str, Any]  )       ->     List[str]:

       """simple docstring"""






       Accelerator()
       lowerCamelCase_							=   DataLoaderDispatcher(range(16  )						,							batch_size=4  )
       for idx, _ in enumerate(A_  ):
             self.assertEqual(dataloader.end_of_dataloader						,							idx == 3  )

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









import math

def       _SCREAMING_SNAKE_CASE      (							lowercase     :	int					):


      '''simple docstring'''



      lowerCamelCase_							=   0
      lowerCamelCase_							=   0
      while num > 0:
            lowerCamelCase_							=   num % 8
            lowerCamelCase_							=   octal + (remainder * math.floor(math.pow(10     ,							lowercase					)					))
            counter += 1
            lowerCamelCase_							=   math.floor(num / 8					)  # basically /= 8 without remainder if any
            # This formatting removes trailing '.0' from `octal`.
      return f"""0o{int(lowercase					)}"""

def       _SCREAMING_SNAKE_CASE      (							):


      '''simple docstring'''



      print('\n2 in octal is:'					)
      print(decimal_to_octal(2					)					)  # = 2
      print('\n8 in octal is:'					)
      print(decimal_to_octal(8					)					)  # = 10
      print('\n65 in octal is:'					)
      print(decimal_to_octal(65					)					)  # = 101
      print('\n216 in octal is:'					)
      print(decimal_to_octal(2_16					)					)  # = 330
      print('\n512 in octal is:'					)
      print(decimal_to_octal(5_12					)					)  # = 1000
      print('\n'					)


if __name__ == "__main__":
 main()



