


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



		"""simple docstring"""

		def __init__( self		,      A=0.01		,      A=10_00	)  -> str:







					'''simple docstring'''







					lowerCamelCase    		=  p_stop
					lowerCamelCase    		=  max_length







		def __iter__( self	)  -> str:







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



		"""simple docstring"""

		def 					__A   ( self		,      A		,      A		,      A=False		,      A=True	)  -> Tuple:







					'''simple docstring'''







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







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







					'''simple docstring'''







					lowerCamelCase    		=  BatchSampler(range(24	)		,      batch_size=3		,      drop_last=A	)
					lowerCamelCase    		=  [
					    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
					    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
					]
					self.check_batch_sampler_shards(A		,      A	)

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

					# Check the shards when the dataset is a round multiple of batch size but not total batch size.
					lowerCamelCase    		=  BatchSampler(range(21	)		,      batch_size=3		,      drop_last=A	)
					lowerCamelCase    		=  [
					    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
					    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [0, 1, 2]],
					]
					self.check_batch_sampler_shards(A		,      A	)

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

					# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
					# num_processes batch.
					lowerCamelCase    		=  BatchSampler(range(22	)		,      batch_size=3		,      drop_last=A	)
					lowerCamelCase    		=  [
					    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
					    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 0, 1]],
					]
					self.check_batch_sampler_shards(A		,      A	)

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

					# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
					# num_processes batch.
					lowerCamelCase    		=  BatchSampler(range(20	)		,      batch_size=3		,      drop_last=A	)
					lowerCamelCase    		=  [
					    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 0]],
					    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [1, 2, 3]],
					]
					self.check_batch_sampler_shards(A		,      A	)

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

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

					lowerCamelCase    		=  BatchSampler(range(2	)		,      batch_size=3		,      drop_last=A	)
					lowerCamelCase    		=  [[], []]
					self.check_batch_sampler_shards(A		,      A	)







		def 					__A   ( self	)  -> str:







					'''simple docstring'''







					lowerCamelCase    		=  BatchSampler(range(24	)		,      batch_size=4		,      drop_last=A	)
					lowerCamelCase    		=  [
					    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
					    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
					]
					self.check_batch_sampler_shards(A		,      A		,      split_batches=A	)

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

					# Check the shards when the dataset is not a round multiple of batch size.
					lowerCamelCase    		=  BatchSampler(range(22	)		,      batch_size=4		,      drop_last=A	)
					lowerCamelCase    		=  [
					    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
					    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [0, 1]],
					]
					self.check_batch_sampler_shards(A		,      A		,      split_batches=A	)

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

					# Check the shards when the dataset is not a round multiple of batch size or num_processes.
					lowerCamelCase    		=  BatchSampler(range(21	)		,      batch_size=4		,      drop_last=A	)
					lowerCamelCase    		=  [
					    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 0]],
					    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [1, 2]],
					]
					self.check_batch_sampler_shards(A		,      A		,      split_batches=A	)

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

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

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







		def 					__A   ( self	)  -> Any:







					'''simple docstring'''







					lowerCamelCase    		=  BatchSampler(range(24	)		,      batch_size=3		,      drop_last=A	)
					lowerCamelCase    		=  [
					    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
					    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
					]
					self.check_batch_sampler_shards(A		,      A		,      even_batches=A	)

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

					# Check the shards when the dataset is a round multiple of batch size but not total batch size.
					lowerCamelCase    		=  BatchSampler(range(21	)		,      batch_size=3		,      drop_last=A	)
					lowerCamelCase    		=  [
					    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
					    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
					]
					self.check_batch_sampler_shards(A		,      A		,      even_batches=A	)

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

					# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
					# num_processes batch.
					lowerCamelCase    		=  BatchSampler(range(22	)		,      batch_size=3		,      drop_last=A	)
					lowerCamelCase    		=  [
					    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
					    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21]],
					]
					self.check_batch_sampler_shards(A		,      A		,      even_batches=A	)

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

					# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
					# num_processes batch.
					lowerCamelCase    		=  BatchSampler(range(20	)		,      batch_size=3		,      drop_last=A	)
					lowerCamelCase    		=  [
					    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19]],
					    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
					]
					self.check_batch_sampler_shards(A		,      A		,      even_batches=A	)

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

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

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







		def 					__A   ( self	)  -> str:







					'''simple docstring'''







					lowerCamelCase    		=  BatchSampler(range(24	)		,      batch_size=4		,      drop_last=A	)
					lowerCamelCase    		=  [
					    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
					    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
					]
					self.check_batch_sampler_shards(A		,      A		,      split_batches=A		,      even_batches=A	)

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

					# Check the shards when the dataset is not a round multiple of batch size.
					lowerCamelCase    		=  BatchSampler(range(22	)		,      batch_size=4		,      drop_last=A	)
					lowerCamelCase    		=  [
					    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
					    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
					]
					self.check_batch_sampler_shards(A		,      A		,      split_batches=A		,      even_batches=A	)

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

					# Check the shards when the dataset is not a round multiple of batch size or num_processes.
					lowerCamelCase    		=  BatchSampler(range(21	)		,      batch_size=4		,      drop_last=A	)
					lowerCamelCase    		=  [
					    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20]],
					    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
					]
					self.check_batch_sampler_shards(A		,      A		,      split_batches=A		,      even_batches=A	)

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

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

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







		def 					__A   ( self	)  -> int:







					'''simple docstring'''







					lowerCamelCase    		=  [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
					lowerCamelCase    		=  [BatchSamplerShard(A		,      2		,      A		,      even_batches=A	) for i in range(2	)]

					self.assertEqual(len(batch_sampler_shards[0]	)		,      3	)
					self.assertEqual(len(batch_sampler_shards[1]	)		,      2	)

					self.assertListEqual(list(batch_sampler_shards[0]	)		,      [[0, 1, 2], [5, 6, 7, 8], [12, 13]]	)
					self.assertListEqual(list(batch_sampler_shards[1]	)		,      [[3, 4], [9, 10, 11]]	)







		def 					__A   ( self		,      A		,      A		,      A		,      A=False		,      A=2		,      A=False	)  -> List[Any]:







					'''simple docstring'''







					random.seed(A	)
					lowerCamelCase    		=  list(A	)

					lowerCamelCase    		=  [
					    IterableDatasetShard(
					        A		,      batch_size=A		,      drop_last=A		,      num_processes=A		,      process_index=A		,      split_batches=A		,      )
					    for i in range(A	)
					]
					lowerCamelCase    		=  []
					for iterable_dataset_shard in iterable_dataset_shards:
								# Since our random iterable dataset will be... random... we need to use a seed to get reproducible results.
								random.seed(A	)
								iterable_dataset_lists.append(list(A	)	)

					lowerCamelCase    		=  batch_size // num_processes if split_batches else batch_size
					# All iterable dataset shard should have the same length, a round multiple of shard_batch_size
					lowerCamelCase    		=  iterable_dataset_lists[0]
					for l in iterable_dataset_lists[1:]:
								self.assertEqual(len(A	)		,      len(A	)	)
								self.assertTrue(len(A	) % shard_batch_size == 0	)

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

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







		def 					__A   ( self	)  -> str:







					'''simple docstring'''







					lowerCamelCase    		=  42
					lowerCamelCase    		=  RandomIterableDataset()

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

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

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







		def 					__A   ( self	)  -> str:







					'''simple docstring'''







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







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







					'''simple docstring'''







					lowerCamelCase    		=  SkipDataLoader(list(range(16	)	)		,      batch_size=4		,      skip_batches=2	)
					self.assertListEqual([t.tolist() for t in dataloader]		,      [[8, 9, 10, 11], [12, 13, 14, 15]]	)







		def 					__A   ( self	)  -> str:







					'''simple docstring'''







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







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







					'''simple docstring'''







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

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







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







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





UpperCAmelCase   :    Any     =		[0, 2, 4, 6, 8]
UpperCAmelCase   :    Optional[Any]     =		[1, 3, 5, 7, 9]






def 	__lowerCamelCase						(							lowerCamelCase__		:						int  ,							lowerCamelCase__		:						int  ,							lowerCamelCase__		:						list[int]  ,							lowerCamelCase__		:						int      ):





			'''simple docstring'''




			if remaining_length == 0:
						if digits[0] == 0 or digits[-1] == 0:
									return 0

						for i in range(length // 2 - 1  ,							-1  ,							-1      ):
									remainder += digits[i] + digits[length - i - 1]

									if remainder % 2 == 0:
												return 0

									remainder //= 10

						return 1

			if remaining_length == 1:
						if remainder % 2 == 0:
									return 0

						lowerCamelCase    		=  0
						for digit in range(10      ):
									lowerCamelCase    		=  digit
									result += reversible_numbers(
									    0  ,							(remainder + 2 * digit) // 10  ,							lowerCamelCase__  ,							lowerCamelCase__      )
						return result

			lowerCamelCase    		=  0
			for digita in range(10      ):
						lowerCamelCase    		=  digita

						if (remainder + digita) % 2 == 0:
									lowerCamelCase    		=  ODD_DIGITS
						else:
									lowerCamelCase    		=  EVEN_DIGITS

						for digita in other_parity_digits:
									lowerCamelCase    		=  digita
									result += reversible_numbers(
									    remaining_length - 2  ,							(remainder + digita + digita) // 10  ,							lowerCamelCase__  ,							lowerCamelCase__  ,							)
			return result






def 	__lowerCamelCase						(							lowerCamelCase__		:						int = 9      ):





			'''simple docstring'''




			lowerCamelCase    		=  0
			for length in range(1  ,							max_power + 1      ):
						result += reversible_numbers(lowerCamelCase__  ,							0  ,							[0] * length  ,							lowerCamelCase__      )
			return result


if __name__ == "__main__":
		print(f"""{solution() = }""")

