




'''simple docstring'''


class  __lowercase    :







   def __init__(						self   ,				UpperCamelCase   )					->     None:
     __a								=       len(UpperCamelCase   )
     __a								=       [0] * len_array

     if len_array > 0:
       __a								=       array[0]

     for i in range(1   ,				UpperCamelCase   ):
       __a								=       self.prefix_sum[i - 1] + array[i]







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   )					->     int:

     if start == 0:
       return self.prefix_sum[end]

     return self.prefix_sum[end] - self.prefix_sum[start - 1]







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase   )					->     bool:
     __a								=       {0}
     for sum_item in self.prefix_sum:
       if sum_item - target_sum in sums:
         return True

       sums.add(UpperCamelCase   )

     return False


if __name__ == "__main__":
  import doctest

  doctest.testmod()





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 							lowerCAmelCase__   ( __lowercase						):

						def __init__(      self							,				a=0.01							,				a=10_00					)       ->   Optional[int]:


								'''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 							lowerCAmelCase__   ( unittest.TestCase						):

						def 			A_					(      self							,				a							,				a							,				a=False							,				a=True					)       ->   Optional[Any]:


								'''simple docstring'''

								_UpperCamelCase         =						[
								    BatchSamplerShard(a							,				2							,				a							,				split_batches=a							,				even_batches=a					)
								    for i in range(2					)
								]
								_UpperCamelCase         =						[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					)       ->   Any:


								'''simple docstring'''

								_UpperCamelCase         =						BatchSampler(range(24					)							,				batch_size=3							,				drop_last=a					)
								_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(a							,				a					)

								_UpperCamelCase         =						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.
								_UpperCamelCase         =						BatchSampler(range(21					)							,				batch_size=3							,				drop_last=a					)
								_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(a							,				a					)

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

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

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

								_UpperCamelCase         =						BatchSampler(range(2					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[[], []]
								self.check_batch_sampler_shards(a							,				a					)

						def 			A_					(      self					)       ->   Tuple:


								'''simple docstring'''

								_UpperCamelCase         =						BatchSampler(range(24					)							,				batch_size=4							,				drop_last=a					)
								_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(a							,				a							,				split_batches=a					)

								_UpperCamelCase         =						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.
								_UpperCamelCase         =						BatchSampler(range(22					)							,				batch_size=4							,				drop_last=a					)
								_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(a							,				a							,				split_batches=a					)

								_UpperCamelCase         =						BatchSampler(range(22					)							,				batch_size=4							,				drop_last=a					)
								_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(a							,				a							,				split_batches=a					)

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

								_UpperCamelCase         =						BatchSampler(range(21					)							,				batch_size=4							,				drop_last=a					)
								_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(a							,				a							,				split_batches=a					)

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

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

						def 			A_					(      self					)       ->   str:


								'''simple docstring'''

								_UpperCamelCase         =						BatchSampler(range(24					)							,				batch_size=3							,				drop_last=a					)
								_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(a							,				a							,				even_batches=a					)

								_UpperCamelCase         =						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.
								_UpperCamelCase         =						BatchSampler(range(21					)							,				batch_size=3							,				drop_last=a					)
								_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(a							,				a							,				even_batches=a					)

								_UpperCamelCase         =						BatchSampler(range(21					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[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.
								_UpperCamelCase         =						BatchSampler(range(22					)							,				batch_size=3							,				drop_last=a					)
								_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(a							,				a							,				even_batches=a					)

								_UpperCamelCase         =						BatchSampler(range(22					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[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.
								_UpperCamelCase         =						BatchSampler(range(20					)							,				batch_size=3							,				drop_last=a					)
								_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(a							,				a							,				even_batches=a					)

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

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

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


								'''simple docstring'''

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

								_UpperCamelCase         =						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.
								_UpperCamelCase         =						BatchSampler(range(22					)							,				batch_size=4							,				drop_last=a					)
								_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(a							,				a							,				split_batches=a							,				even_batches=a					)

								_UpperCamelCase         =						BatchSampler(range(22					)							,				batch_size=4							,				drop_last=a					)
								_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(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.
								_UpperCamelCase         =						BatchSampler(range(21					)							,				batch_size=4							,				drop_last=a					)
								_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(a							,				a							,				split_batches=a							,				even_batches=a					)

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

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

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

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


								'''simple docstring'''

								_UpperCamelCase         =						[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
								_UpperCamelCase         =						[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							,				a							,				a							,				a							,				a=False							,				a=2							,				a=False					)       ->   Any:


								'''simple docstring'''

								random.seed(a					)
								_UpperCamelCase         =						list(a					)

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

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

								_UpperCamelCase         =						[]
								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					)       ->   int:


								'''simple docstring'''

								_UpperCamelCase         =						42
								_UpperCamelCase         =						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
								_UpperCamelCase         =						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					)       ->   Any:


								'''simple docstring'''

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

						def 			A_					(      self					)       ->   Dict:


								'''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 			A_					(      self					)       ->   Tuple:


								'''simple docstring'''

								_UpperCamelCase         =						DataLoader(list(range(16					)					)							,				batch_size=4					)
								_UpperCamelCase         =						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					)       ->   List[str]:


								'''simple docstring'''

								_UpperCamelCase         =						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					)       ->   Optional[Any]:


								'''simple docstring'''

								Accelerator()
								_UpperCamelCase         =						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					)



