




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




						"""simple docstring"""





						def __init__(				self    ,				A=0.01    ,				A=1_0_0_0     )		->    List[str]:
									snake_case				:    List[Any]           =	p_stop
									snake_case				:    List[Any]           =	max_length

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




class 						__lowercase		(unittest.TestCase							):




						"""simple docstring"""





						def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A=False    ,				A=True     )		->    Union[str, Any]:
									snake_case				:    List[str]           =	[
									    BatchSamplerShard(lowercase_    ,				2    ,				lowercase_    ,				split_batches=lowercase_    ,				even_batches=lowercase_     )
									    for i in range(2     )
									]
									snake_case				:    Dict           =	[list(lowercase_     ) for batch_sampler_shard in batch_sampler_shards]
									if not split_batches:
												self.assertListEqual([len(lowercase_     ) for shard in batch_sampler_shards]    ,				[len(lowercase_     ) for e in expected]     )
									self.assertListEqual(lowercase_    ,				lowercase_     )

						def      UpperCAmelCase				(				self     )		->    int:
									# Check the shards when the dataset is a round multiple of total batch size.
									snake_case				:    Optional[Any]           =	BatchSampler(range(2_4     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    int           =	[
									    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
									    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1, 2_2, 2_3]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_     )

									snake_case				:    Any           =	BatchSampler(range(2_4     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									# Expected shouldn't change
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_     )

									# Check the shards when the dataset is a round multiple of batch size but not total batch size.
									snake_case				:    Union[str, Any]           =	BatchSampler(range(2_1     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    Any           =	[
									    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
									    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [0, 1, 2]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_     )

									snake_case				:    int           =	BatchSampler(range(2_1     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    Dict           =	[
									    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
									    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_     )

									# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
									# num_processes batch.
									snake_case				:    List[Any]           =	BatchSampler(range(2_2     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    List[Any]           =	[
									    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
									    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1, 0, 1]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_     )

									snake_case				:    Optional[int]           =	BatchSampler(range(2_2     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    Optional[int]           =	[
									    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
									    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_     )

									# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
									# num_processes batch.
									snake_case				:    Optional[Any]           =	BatchSampler(range(2_0     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    List[str]           =	[
									    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 0]],
									    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [1, 2, 3]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_     )

									snake_case				:    Union[str, Any]           =	BatchSampler(range(2_0     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    List[str]           =	[
									    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
									    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_     )

									# Check the shards when the dataset is very small.
									snake_case				:    int           =	BatchSampler(range(2     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    Optional[Any]           =	[[[0, 1, 0]], [[1, 0, 1]]]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_     )

									snake_case				:    Optional[int]           =	BatchSampler(range(2     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    Any           =	[[], []]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_     )

						def      UpperCAmelCase				(				self     )		->    List[str]:
									# Check the shards when the dataset is a round multiple of batch size.
									snake_case				:    Optional[Any]           =	BatchSampler(range(2_4     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									snake_case				:    Union[str, Any]           =	[
									    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
									    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [2_2, 2_3]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				split_batches=lowercase_     )

									snake_case				:    Union[str, Any]           =	BatchSampler(range(2_4     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									# Expected shouldn't change
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				split_batches=lowercase_     )

									# Check the shards when the dataset is not a round multiple of batch size.
									snake_case				:    Union[str, Any]           =	BatchSampler(range(2_2     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									snake_case				:    Any           =	[
									    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
									    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [0, 1]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				split_batches=lowercase_     )

									snake_case				:    Dict           =	BatchSampler(range(2_2     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									snake_case				:    List[Any]           =	[
									    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
									    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				split_batches=lowercase_     )

									# Check the shards when the dataset is not a round multiple of batch size or num_processes.
									snake_case				:    str           =	BatchSampler(range(2_1     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									snake_case				:    List[str]           =	[
									    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 0]],
									    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [1, 2]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				split_batches=lowercase_     )

									snake_case				:    str           =	BatchSampler(range(2_1     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									snake_case				:    str           =	[
									    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
									    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				split_batches=lowercase_     )

									# Check the shards when the dataset is very small.
									snake_case				:    int           =	BatchSampler(range(2     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									snake_case				:    Tuple           =	[[[0, 1]], [[0, 1]]]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				split_batches=lowercase_     )

									snake_case				:    List[str]           =	BatchSampler(range(2     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									snake_case				:    List[str]           =	[[], []]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				split_batches=lowercase_     )

						def      UpperCAmelCase				(				self     )		->    Dict:
									# Check the shards when the dataset is a round multiple of total batch size.
									snake_case				:    List[str]           =	BatchSampler(range(2_4     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    Optional[Any]           =	[
									    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
									    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1, 2_2, 2_3]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				even_batches=lowercase_     )

									snake_case				:    Tuple           =	BatchSampler(range(2_4     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									# Expected shouldn't change
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				even_batches=lowercase_     )

									# Check the shards when the dataset is a round multiple of batch size but not total batch size.
									snake_case				:    Optional[Any]           =	BatchSampler(range(2_1     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    List[Any]           =	[
									    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
									    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				even_batches=lowercase_     )

									snake_case				:    Tuple           =	BatchSampler(range(2_1     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    Any           =	[
									    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
									    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				even_batches=lowercase_     )

									# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
									# num_processes batch.
									snake_case				:    Dict           =	BatchSampler(range(2_2     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    List[str]           =	[
									    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
									    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				even_batches=lowercase_     )

									snake_case				:    Union[str, Any]           =	BatchSampler(range(2_2     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    List[Any]           =	[
									    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
									    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				even_batches=lowercase_     )

									# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
									# num_processes batch.
									snake_case				:    Optional[int]           =	BatchSampler(range(2_0     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    int           =	[
									    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9]],
									    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				even_batches=lowercase_     )

									snake_case				:    Dict           =	BatchSampler(range(2_0     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    Optional[Any]           =	[
									    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
									    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				even_batches=lowercase_     )

									# Check the shards when the dataset is very small.
									snake_case				:    Tuple           =	BatchSampler(range(2     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    List[str]           =	[[[0, 1]], []]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				even_batches=lowercase_     )

									snake_case				:    Optional[int]           =	BatchSampler(range(2     )    ,				batch_size=3    ,				drop_last=lowercase_     )
									snake_case				:    Any           =	[[], []]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				even_batches=lowercase_     )

						def      UpperCAmelCase				(				self     )		->    Union[str, Any]:
									# Check the shards when the dataset is a round multiple of batch size.
									snake_case				:    Tuple           =	BatchSampler(range(2_4     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									snake_case				:    Union[str, Any]           =	[
									    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
									    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [2_2, 2_3]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				split_batches=lowercase_    ,				even_batches=lowercase_     )

									snake_case				:    int           =	BatchSampler(range(2_4     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									# Expected shouldn't change
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				split_batches=lowercase_    ,				even_batches=lowercase_     )

									# Check the shards when the dataset is not a round multiple of batch size.
									snake_case				:    Dict           =	BatchSampler(range(2_2     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									snake_case				:    Any           =	[
									    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
									    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				split_batches=lowercase_    ,				even_batches=lowercase_     )

									snake_case				:    Optional[int]           =	BatchSampler(range(2_2     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									snake_case				:    Optional[int]           =	[
									    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
									    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				split_batches=lowercase_    ,				even_batches=lowercase_     )

									# Check the shards when the dataset is not a round multiple of batch size or num_processes.
									snake_case				:    Any           =	BatchSampler(range(2_1     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									snake_case				:    int           =	[
									    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0]],
									    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				split_batches=lowercase_    ,				even_batches=lowercase_     )

									snake_case				:    Union[str, Any]           =	BatchSampler(range(2_1     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									snake_case				:    int           =	[
									    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
									    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
									]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				split_batches=lowercase_    ,				even_batches=lowercase_     )

									# Check the shards when the dataset is very small.
									snake_case				:    Optional[Any]           =	BatchSampler(range(2     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									snake_case				:    int           =	[[[0, 1]], []]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				split_batches=lowercase_    ,				even_batches=lowercase_     )

									snake_case				:    str           =	BatchSampler(range(2     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									snake_case				:    str           =	[[], []]
									self.check_batch_sampler_shards(lowercase_    ,				lowercase_    ,				split_batches=lowercase_    ,				even_batches=lowercase_     )

						def      UpperCAmelCase				(				self     )		->    str:
									snake_case				:    List[str]           =	[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 1_0, 1_1], [1_2, 1_3]]
									snake_case				:    str           =	[BatchSamplerShard(lowercase_    ,				2    ,				lowercase_    ,				even_batches=lowercase_     ) 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], [1_2, 1_3]]     )
									self.assertListEqual(list(batch_sampler_shards[1]     )    ,				[[3, 4], [9, 1_0, 1_1]]     )

						def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A    ,				A=False    ,				A=2    ,				A=False     )		->    str:
									random.seed(lowercase_     )
									snake_case				:    List[str]           =	list(lowercase_     )

									snake_case				:    Optional[Any]           =	[
									    IterableDatasetShard(
									        lowercase_    ,				batch_size=lowercase_    ,				drop_last=lowercase_    ,				num_processes=lowercase_    ,				process_index=lowercase_    ,				split_batches=lowercase_    ,				)
									    for i in range(lowercase_     )
									]
									snake_case				:    List[str]           =	[]
									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(lowercase_     )
												iterable_dataset_lists.append(list(lowercase_     )     )

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

									snake_case				:    int           =	[]
									for idx in range(0    ,				len(lowercase_     )    ,				lowercase_     ):
												for l in iterable_dataset_lists:
															observed += l[idx : idx + shard_batch_size]

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

						def      UpperCAmelCase				(				self     )		->    str:
									snake_case				:    Tuple           =	4_2
									snake_case				:    Dict           =	RandomIterableDataset()

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

									# Edge case with a very small dataset
									snake_case				:    Tuple           =	RandomIterableDataset(max_length=2     )

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

						def      UpperCAmelCase				(				self     )		->    List[Any]:
									snake_case				:    List[Any]           =	BatchSampler(range(1_6     )    ,				batch_size=4    ,				drop_last=lowercase_     )
									snake_case				:    int           =	SkipBatchSampler(lowercase_    ,				2     )
									self.assertListEqual(list(lowercase_     )    ,				[[8, 9, 1_0, 1_1], [1_2, 1_3, 1_4, 1_5]]     )

						def      UpperCAmelCase				(				self     )		->    int:
									snake_case				:    int           =	SkipDataLoader(list(range(1_6     )     )    ,				batch_size=4    ,				skip_batches=2     )
									self.assertListEqual([t.tolist() for t in dataloader]    ,				[[8, 9, 1_0, 1_1], [1_2, 1_3, 1_4, 1_5]]     )

						def      UpperCAmelCase				(				self     )		->    Any:
									snake_case				:    Optional[Any]           =	DataLoader(list(range(1_6     )     )    ,				batch_size=4     )
									snake_case				:    int           =	skip_first_batches(lowercase_    ,				num_batches=2     )
									self.assertListEqual([t.tolist() for t in new_dataloader]    ,				[[8, 9, 1_0, 1_1], [1_2, 1_3, 1_4, 1_5]]     )

						def      UpperCAmelCase				(				self     )		->    Union[str, Any]:
									snake_case				:    int           =	DataLoaderShard(list(range(1_6     )     )    ,				batch_size=4     )
									for idx, _ in enumerate(lowercase_     ):
												self.assertEqual(dataloader.end_of_dataloader    ,				idx == 3     )

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

						def      UpperCAmelCase				(				self     )		->    List[str]:
									Accelerator()
									snake_case				:    Tuple           =	DataLoaderDispatcher(range(1_6     )    ,				batch_size=4     )
									for idx, _ in enumerate(lowercase_     ):
												self.assertEqual(dataloader.end_of_dataloader    ,				idx == 3     )

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











def 						SCREAMING_SNAKE_CASE__ (   lowercase						,lowercase					)  ->       int:
   return int((input_a, input_a).count(1					) != 0					)




def 						SCREAMING_SNAKE_CASE__ (   )  ->       None:
   assert or_gate(0						,0					) == 0
   assert or_gate(0						,1					) == 1
   assert or_gate(1						,0					) == 1
   assert or_gate(1						,1					) == 1


if __name__ == "__main__":
       print(or_gate(0, 1))
       print(or_gate(1, 0))
       print(or_gate(0, 0))
       print(or_gate(1, 1))





