


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     lowercase__  (					__SCREAMING_SNAKE_CASE		):






  def __init__(					self      :					int		,						_lowercase      :					Union[str, Any]=0.0_1		,						_lowercase      :					Tuple=10_00 ):




         """simple docstring"""







         UpperCAmelCase__										=     p_stop
         UpperCAmelCase__										=     max_length






  def __iter__(					self      :					Optional[Any] ):




         """simple docstring"""







         UpperCAmelCase__										=     0
         UpperCAmelCase__										=     False
         while not stop and count < self.max_length:
                yield count
                count += 1
                UpperCAmelCase__										=     random.random() < self.p_stop




class     lowercase__  (					unittest.TestCase		):






  def 			_UpperCAmelCase  (					self      :					List[str]		,						_lowercase      :					List[str]		,						_lowercase      :					int		,						_lowercase      :					List[str]=False		,						_lowercase      :					List[Any]=True ):




         """simple docstring"""







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






  def 			_UpperCAmelCase  (					self      :					Tuple ):




         """simple docstring"""







         UpperCAmelCase__										=     BatchSampler(range(24 )		,						batch_size=3		,						drop_last=_lowercase )
         UpperCAmelCase__										=     [
             [[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(_lowercase		,						_lowercase )

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

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

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

         # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
         # num_processes batch.
         UpperCAmelCase__										=     BatchSampler(range(22 )		,						batch_size=3		,						drop_last=_lowercase )
         UpperCAmelCase__										=     [
             [[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(_lowercase		,						_lowercase )

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

         # 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__										=     BatchSampler(range(20 )		,						batch_size=3		,						drop_last=_lowercase )
         UpperCAmelCase__										=     [
             [[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(_lowercase		,						_lowercase )

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

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

         UpperCAmelCase__										=     BatchSampler(range(2 )		,						batch_size=3		,						drop_last=_lowercase )
         UpperCAmelCase__										=     [[], []]
         self.check_batch_sampler_shards(_lowercase		,						_lowercase )






  def 			_UpperCAmelCase  (					self      :					int ):




         """simple docstring"""







         UpperCAmelCase__										=     BatchSampler(range(24 )		,						batch_size=4		,						drop_last=_lowercase )
         UpperCAmelCase__										=     [
             [[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(_lowercase		,						_lowercase		,						split_batches=_lowercase )

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

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

         UpperCAmelCase__										=     BatchSampler(range(22 )		,						batch_size=4		,						drop_last=_lowercase )
         UpperCAmelCase__										=     [
             [[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(_lowercase		,						_lowercase		,						split_batches=_lowercase )

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

         UpperCAmelCase__										=     BatchSampler(range(21 )		,						batch_size=4		,						drop_last=_lowercase )
         UpperCAmelCase__										=     [
             [[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(_lowercase		,						_lowercase		,						split_batches=_lowercase )

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

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






  def 			_UpperCAmelCase  (					self      :					Optional[int] ):




         """simple docstring"""







         UpperCAmelCase__										=     BatchSampler(range(24 )		,						batch_size=3		,						drop_last=_lowercase )
         UpperCAmelCase__										=     [
             [[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(_lowercase		,						_lowercase		,						even_batches=_lowercase )

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

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

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

         # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
         # num_processes batch.
         UpperCAmelCase__										=     BatchSampler(range(22 )		,						batch_size=3		,						drop_last=_lowercase )
         UpperCAmelCase__										=     [
             [[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(_lowercase		,						_lowercase		,						even_batches=_lowercase )

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

         # 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__										=     BatchSampler(range(20 )		,						batch_size=3		,						drop_last=_lowercase )
         UpperCAmelCase__										=     [
             [[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(_lowercase		,						_lowercase		,						even_batches=_lowercase )

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

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

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






  def 			_UpperCAmelCase  (					self      :					Optional[Any] ):




         """simple docstring"""







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

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

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

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

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

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

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

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






  def 			_UpperCAmelCase  (					self      :					Optional[Any] ):




         """simple docstring"""







         UpperCAmelCase__										=     [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
         UpperCAmelCase__										=     [BatchSamplerShard(_lowercase		,						2		,						_lowercase		,						even_batches=_lowercase ) 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 			_UpperCAmelCase  (					self      :					str		,						_lowercase      :					Optional[int]		,						_lowercase      :					str		,						_lowercase      :					Dict		,						_lowercase      :					List[str]=False		,						_lowercase      :					Tuple=2		,						_lowercase      :					Dict=False ):




         """simple docstring"""







         random.seed(_lowercase )
         UpperCAmelCase__										=     list(_lowercase )

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

         UpperCAmelCase__										=     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__										=     iterable_dataset_lists[0]
         for l in iterable_dataset_lists[1:]:
                self.assertEqual(len(_lowercase )		,						len(_lowercase ) )
                self.assertTrue(len(_lowercase ) % shard_batch_size == 0 )

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

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






  def 			_UpperCAmelCase  (					self      :					List[str] ):




         """simple docstring"""







         UpperCAmelCase__										=     42
         UpperCAmelCase__										=     RandomIterableDataset()

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

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

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






  def 			_UpperCAmelCase  (					self      :					str ):




         """simple docstring"""







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






  def 			_UpperCAmelCase  (					self      :					str ):




         """simple docstring"""







         UpperCAmelCase__										=     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 			_UpperCAmelCase  (					self      :					List[Any] ):




         """simple docstring"""







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






  def 			_UpperCAmelCase  (					self      :					Optional[Any] ):




         """simple docstring"""







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

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






  def 			_UpperCAmelCase  (					self      :					int ):




         """simple docstring"""







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

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









from .imports import is_tqdm_available


if is_tqdm_available():
    from tqdm.auto import tqdm as _tqdm

from ..state import PartialState





def __UpperCAmelCase       (		__A = True ,						*__A ,						**__A			)						->					Any:



       '''simple docstring'''





       if not is_tqdm_available():
              raise ImportError("Accelerate's `tqdm` module requires `tqdm` to be installed. Please run `pip install tqdm`."			)
       UpperCAmelCase__										=     False
       if main_process_only:
              UpperCAmelCase__										=     PartialState().local_process_index == 0
       return _tqdm(*__A ,						**__A ,						disable=__A			)





