



'''simple docstring'''





import random
import unittest

from torch.utils.data import BatchSampler, DataLoader, IterableDataset

from accelerate import Accelerator
from accelerate.data_loader import (
    BatchSamplerShard,
    DataLoaderDispatcher,
    DataLoaderShard,
    IterableDatasetShard,
    SkipBatchSampler,
    SkipDataLoader,
    skip_first_batches,
)


class 				snake_case (UpperCamelCase   ):



       def __init__(				self     ,UpperCAmelCase_=0.01     ,UpperCAmelCase_=1_000      )  ->							Tuple:
           lowercase__		       =						p_stop
           lowercase__		       =						max_length







       def __iter__(				self      )  ->							List[str]:
           lowercase__		       =						0
           lowercase__		       =						False
           while not stop and count < self.max_length:
               yield count
               count += 1
               lowercase__		       =						random.random() < self.p_stop





class 				snake_case (unittest.TestCase   ):



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



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

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

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

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

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

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

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

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

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

           lowercase__		       =						BatchSampler(range(2      )     ,batch_size=3     ,drop_last=UpperCAmelCase_      )
           lowercase__		       =						[[], []]
           self.check_batch_sampler_shards(UpperCAmelCase_     ,UpperCAmelCase_      )



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

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

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

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

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

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

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

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



       def 							_a      (				self      )  ->							Dict:
           # Check the shards when the dataset is a round multiple of total batch size.
           lowercase__		       =						BatchSampler(range(24      )     ,batch_size=3     ,drop_last=UpperCAmelCase_      )
           lowercase__		       =						[
               [[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(UpperCAmelCase_     ,UpperCAmelCase_     ,even_batches=UpperCAmelCase_      )

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

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

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

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

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

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

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

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

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



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

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

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

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

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

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

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

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



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

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

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

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

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



       def 							_a      (				self      )  ->							List[Any]:
           lowercase__		       =						42
           lowercase__		       =						RandomIterableDataset()

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

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

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



       def 							_a      (				self      )  ->							List[str]:
           lowercase__		       =						BatchSampler(range(16      )     ,batch_size=4     ,drop_last=UpperCAmelCase_      )
           lowercase__		       =						SkipBatchSampler(UpperCAmelCase_     ,2      )
           self.assertListEqual(list(UpperCAmelCase_      )     ,[[8, 9, 10, 11], [12, 13, 14, 15]]      )



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



       def 							_a      (				self      )  ->							Tuple:
           lowercase__		       =						DataLoaderShard(list(range(16      )      )     ,batch_size=4      )
           for idx, _ in enumerate(UpperCAmelCase_      ):
               self.assertEqual(dataloader.end_of_dataloader     ,idx == 3      )

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







       def 							_a      (				self      )  ->							List[Any]:
           Accelerator()
           lowercase__		       =						DataLoaderDispatcher(range(16      )     ,batch_size=4      )
           for idx, _ in enumerate(UpperCAmelCase_      ):
               self.assertEqual(dataloader.end_of_dataloader     ,idx == 3      )

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









'''simple docstring'''





# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
import platform

import numpy as np
import psutil
import torch

from accelerate import __version__ as version
from accelerate.commands.config import default_config_file, load_config_from_file

from ..utils import is_npu_available, is_xpu_available




def        lowerCamelCase  (				_snake_case  :						str=None       ):





    '''simple docstring'''


    if subparsers is not None:
        lowercase__		       =						subparsers.add_parser("env"       )
    else:
        lowercase__		       =						argparse.ArgumentParser("Accelerate env command"       )

    parser.add_argument(
        "--config_file"     ,default=_snake_case     ,help="The config file to use for the default values in the launching script."       )

    if subparsers is not None:
        parser.set_defaults(func=_snake_case       )
    return parser




def        lowerCamelCase  (				_snake_case  :						Tuple       ):





    '''simple docstring'''


    lowercase__		       =						torch.__version__
    lowercase__		       =						torch.cuda.is_available()
    lowercase__		       =						is_xpu_available()
    lowercase__		       =						is_npu_available()

    lowercase__		       =						"Not found"
    # Get the default from the config file.
    if args.config_file is not None or os.path.isfile(_snake_case       ):
        lowercase__		       =						load_config_from_file(args.config_file       ).to_dict()

    lowercase__		       =						{
        "`Accelerate` version": version,
        "Platform": platform.platform(),
        "Python version": platform.python_version(),
        "Numpy version": np.__version__,
        "PyTorch version (GPU?)": f'''{pt_version} ({pt_cuda_available})''',
        "PyTorch XPU available": str(_snake_case       ),
        "PyTorch NPU available": str(_snake_case       ),
        "System RAM": f'''{psutil.virtual_memory().total / 1_024 ** 3:.2f} GB''',
    }
    if pt_cuda_available:
        lowercase__		       =						torch.cuda.get_device_name()

    print("\nCopy-and-paste the text below in your GitHub issue\n"       )
    print("\n".join([f'''- {prop}: {val}''' for prop, val in info.items()]       )       )

    print("- `Accelerate` default config:" if args.config_file is None else "- `Accelerate` config passed:"       )
    lowercase__		       =						(
        "\n".join([f'''\t- {prop}: {val}''' for prop, val in accelerate_config.items()]       )
        if isinstance(_snake_case     ,_snake_case       )
        else f'''\t{accelerate_config}'''
    )
    print(_snake_case       )

    lowercase__		       =						accelerate_config

    return info




def        lowerCamelCase  (				):





    '''simple docstring'''


    lowercase__		       =						env_command_parser()
    lowercase__		       =						parser.parse_args()
    env_command(_snake_case       )
    return 0


if __name__ == "__main__":
 raise SystemExit(main())




