




"""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        __magic_name__				(			UpperCAmelCase__							):






				'''simple docstring'''





				def __init__(			self						,					_a=0.01						,					_a=1_000					):





							"""simple docstring"""
							lowerCamelCase      =   p_stop
							lowerCamelCase      =   max_length





				def __iter__(			self					):





							"""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        __magic_name__				(			unittest.TestCase							):






				'''simple docstring'''





				def   _lowerCAmelCase							(			self						,					_a						,					_a						,					_a=False						,					_a=True					):





							"""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   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							# Check the shards when the dataset is a round multiple of total batch size.
							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   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							# Check the shards when the dataset is a round multiple of batch size.
							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   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							# Check the shards when the dataset is a round multiple of total batch size.
							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   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							# Check the shards when the dataset is a round multiple of batch size.
							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   _lowerCAmelCase							(			self					):





							"""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   _lowerCAmelCase							(			self						,					_a						,					_a						,					_a						,					_a=False						,					_a=2						,					_a=False					):





							"""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   _lowerCAmelCase							(			self					):





							"""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   _lowerCAmelCase							(			self					):





							"""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   _lowerCAmelCase							(			self					):





							"""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   _lowerCAmelCase							(			self					):





							"""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   _lowerCAmelCase							(			self					):





							"""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   _lowerCAmelCase							(			self					):





							"""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					)









"""simple docstring"""




from collections.abc import Sequence



def        a__					(							snake_case__			,					snake_case__ = False							)		->   float:
			if not arr:
						return 0

			lowerCamelCase      =   0 if allow_empty_subarrays else float("""-inf"""							)
			lowerCamelCase      =   0.0
			for num in arr:
						lowerCamelCase      =   max(0 if allow_empty_subarrays else num			,					curr_sum + num							)
						lowerCamelCase      =   max(snake_case__			,					snake_case__							)

			return max_sum


if __name__ == "__main__":
					from doctest import testmod

					testmod()

					lowerCAmelCase				:     Union[str, Any]					=				[-2, 1, -3, 4, -1, 2, 1, -5, 4]
					print(F"""{max_subarray_sum(nums) = }""")



