






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_       (				snake_case__					):




							def __init__(   self     :   Optional[int] ,							SCREAMING_SNAKE_CASE_     :   Optional[int]=0.01 ,							SCREAMING_SNAKE_CASE_     :   List[Any]=1_000	):
											lowerCAmelCase__      			=						p_stop
											lowerCAmelCase__      			=						max_length




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



class  lowerCAmelCase_       (				unittest.TestCase					):




							def 						__snake_case					(   self     :   Union[str, Any] ,							SCREAMING_SNAKE_CASE_     :   Union[str, Any] ,							SCREAMING_SNAKE_CASE_     :   Tuple ,							SCREAMING_SNAKE_CASE_     :   Tuple=False ,							SCREAMING_SNAKE_CASE_     :   int=True	):
											lowerCAmelCase__      			=						[
											    BatchSamplerShard(SCREAMING_SNAKE_CASE_ ,							2 ,							SCREAMING_SNAKE_CASE_ ,							split_batches=SCREAMING_SNAKE_CASE_ ,							even_batches=SCREAMING_SNAKE_CASE_	)
											    for i in range(2	)
											]
											lowerCAmelCase__      			=						[list(SCREAMING_SNAKE_CASE_	) for batch_sampler_shard in batch_sampler_shards]
											if not split_batches:
															self.assertListEqual([len(SCREAMING_SNAKE_CASE_	) for shard in batch_sampler_shards] ,							[len(SCREAMING_SNAKE_CASE_	) for e in expected]	)
											self.assertListEqual(SCREAMING_SNAKE_CASE_ ,							SCREAMING_SNAKE_CASE_	)




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

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

											# 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=SCREAMING_SNAKE_CASE_	)
											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(SCREAMING_SNAKE_CASE_ ,							SCREAMING_SNAKE_CASE_	)

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

											# 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=SCREAMING_SNAKE_CASE_	)
											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(SCREAMING_SNAKE_CASE_ ,							SCREAMING_SNAKE_CASE_	)

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

											# 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=SCREAMING_SNAKE_CASE_	)
											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(SCREAMING_SNAKE_CASE_ ,							SCREAMING_SNAKE_CASE_	)

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

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

											lowerCAmelCase__      			=						BatchSampler(range(2	) ,							batch_size=3 ,							drop_last=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						[[], []]
											self.check_batch_sampler_shards(SCREAMING_SNAKE_CASE_ ,							SCREAMING_SNAKE_CASE_	)




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

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

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

											lowerCAmelCase__      			=						BatchSampler(range(22	) ,							batch_size=4 ,							drop_last=SCREAMING_SNAKE_CASE_	)
											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(SCREAMING_SNAKE_CASE_ ,							SCREAMING_SNAKE_CASE_ ,							split_batches=SCREAMING_SNAKE_CASE_	)

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

											lowerCAmelCase__      			=						BatchSampler(range(21	) ,							batch_size=4 ,							drop_last=SCREAMING_SNAKE_CASE_	)
											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(SCREAMING_SNAKE_CASE_ ,							SCREAMING_SNAKE_CASE_ ,							split_batches=SCREAMING_SNAKE_CASE_	)

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

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




							def 						__snake_case					(   self     :   Union[str, Any]	):
											# Check the shards when the dataset is a round multiple of total batch size.
											lowerCAmelCase__      			=						BatchSampler(range(24	) ,							batch_size=3 ,							drop_last=SCREAMING_SNAKE_CASE_	)
											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(SCREAMING_SNAKE_CASE_ ,							SCREAMING_SNAKE_CASE_ ,							even_batches=SCREAMING_SNAKE_CASE_	)

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

											# 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=SCREAMING_SNAKE_CASE_	)
											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(SCREAMING_SNAKE_CASE_ ,							SCREAMING_SNAKE_CASE_ ,							even_batches=SCREAMING_SNAKE_CASE_	)

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

											# 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=SCREAMING_SNAKE_CASE_	)
											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(SCREAMING_SNAKE_CASE_ ,							SCREAMING_SNAKE_CASE_ ,							even_batches=SCREAMING_SNAKE_CASE_	)

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

											# 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=SCREAMING_SNAKE_CASE_	)
											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(SCREAMING_SNAKE_CASE_ ,							SCREAMING_SNAKE_CASE_ ,							even_batches=SCREAMING_SNAKE_CASE_	)

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

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

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




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

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

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

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

											# 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=SCREAMING_SNAKE_CASE_	)
											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(SCREAMING_SNAKE_CASE_ ,							SCREAMING_SNAKE_CASE_ ,							split_batches=SCREAMING_SNAKE_CASE_ ,							even_batches=SCREAMING_SNAKE_CASE_	)

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

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

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




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

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

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

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

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




							def 						__snake_case					(   self     :   Optional[int]	):
											lowerCAmelCase__      			=						42
											lowerCAmelCase__      			=						RandomIterableDataset()

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

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

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




							def 						__snake_case					(   self     :   List[str]	):
											lowerCAmelCase__      			=						BatchSampler(range(16	) ,							batch_size=4 ,							drop_last=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						SkipBatchSampler(SCREAMING_SNAKE_CASE_ ,							2	)
											self.assertListEqual(list(SCREAMING_SNAKE_CASE_	) ,							[[8, 9, 10, 11], [12, 13, 14, 15]]	)




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




							def 						__snake_case					(   self     :   Tuple	):
											lowerCAmelCase__      			=						DataLoaderShard(list(range(16	)	) ,							batch_size=4	)
											for idx, _ in enumerate(SCREAMING_SNAKE_CASE_	):
															self.assertEqual(dataloader.end_of_dataloader ,							idx == 3	)

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




							def 						__snake_case					(   self     :   str	):
											Accelerator()
											lowerCAmelCase__      			=						DataLoaderDispatcher(range(16	) ,							batch_size=4	)
											for idx, _ in enumerate(SCREAMING_SNAKE_CASE_	):
															self.assertEqual(dataloader.end_of_dataloader ,							idx == 3	)

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













from math import isclose, sqrt





def 	lowerCAmelCase_		(lowercase__						:       float     ,					lowercase__						:       float     ,					lowercase__						:       float       )		->    tuple[float, float, float]:


				'''simple docstring'''




				lowerCAmelCase__      			=						point_y / 4 / point_x
				lowerCAmelCase__      			=						2 * normal_gradient / (1 + normal_gradient * normal_gradient)
				lowerCAmelCase__      			=						(1 - normal_gradient * normal_gradient) / (
				    1 + normal_gradient * normal_gradient
				)
				lowerCAmelCase__      			=						(sa - ca * incoming_gradient) / (ca + sa * incoming_gradient)

				# to find the next point, solve the simultaeneous equations:
				# y^2 + 4x^2 = 100
				# y - b = m * (x - a)
				# ==> A x^2 + B x + C = 0
				lowerCAmelCase__      			=						outgoing_gradient**2 + 4
				lowerCAmelCase__      			=						2 * outgoing_gradient * (point_y - outgoing_gradient * point_x)
				lowerCAmelCase__      			=						(point_y - outgoing_gradient * point_x) ** 2 - 1_00

				lowerCAmelCase__      			=						(
				    -linear_term - sqrt(linear_term**2 - 4 * quadratic_term * constant_term       )
				) / (2 * quadratic_term)
				lowerCAmelCase__      			=						(
				    -linear_term + sqrt(linear_term**2 - 4 * quadratic_term * constant_term       )
				) / (2 * quadratic_term)

				# two solutions, one of which is our input point
				lowerCAmelCase__      			=						x_minus if isclose(lowercase__     ,					lowercase__       ) else x_plus
				lowerCAmelCase__      			=						point_y + outgoing_gradient * (next_x - point_x)

				return next_x, next_y, outgoing_gradient





def 	lowerCAmelCase_		(lowercase__						:       float = 1.4     ,					lowercase__						:       float = -9.6       )		->    int:


				'''simple docstring'''




				lowerCAmelCase__      			=						0
				lowerCAmelCase__      			=						first_x_coord
				lowerCAmelCase__      			=						first_y_coord
				lowerCAmelCase__      			=						(10.1 - point_y) / (0.0 - point_x)

				while not (-0.01 <= point_x <= 0.01 and point_y > 0):
								lowerCAmelCase__			,		lowerCAmelCase__			,		lowerCAmelCase__      			=						next_point(lowercase__     ,					lowercase__     ,					lowercase__       )
								num_reflections += 1

				return num_reflections


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





