

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


							def __init__(			self    ,			_snake_case=0.01    ,			_snake_case=1000)							->				Union[str, Any]:
								UpperCAmelCase_		:      Optional[Any]       	=					p_stop
								UpperCAmelCase_		:      Tuple       	=					max_length





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



class 					lowercase						(  unittest.TestCase ):


							def 						_snake_case							(			self    ,			_snake_case    ,			_snake_case    ,			_snake_case=False    ,			_snake_case=True)							->				Optional[Any]:
								UpperCAmelCase_		:      Optional[Any]       	=					[
								    BatchSamplerShard(_snake_case    ,			2    ,			_snake_case    ,			split_batches=_snake_case    ,			even_batches=_snake_case)
								    for i in range(2)
								]
								UpperCAmelCase_		:      Tuple       	=					[list(_snake_case) for batch_sampler_shard in batch_sampler_shards]
								if not split_batches:
									self.assertListEqual([len(_snake_case) for shard in batch_sampler_shards]    ,			[len(_snake_case) for e in expected])
								self.assertListEqual(_snake_case    ,			_snake_case)


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

								UpperCAmelCase_		:      Tuple       	=					BatchSampler(range(24)    ,			batch_size=3    ,			drop_last=_snake_case)
								# Expected shouldn't change
								self.check_batch_sampler_shards(_snake_case    ,			_snake_case)

								# Check the shards when the dataset is a round multiple of batch size but not total batch size.
								UpperCAmelCase_		:      Tuple       	=					BatchSampler(range(21)    ,			batch_size=3    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      int       	=					[
								    [[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(_snake_case    ,			_snake_case)

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

								# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
								# num_processes batch.
								UpperCAmelCase_		:      Union[str, Any]       	=					BatchSampler(range(22)    ,			batch_size=3    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      Optional[Any]       	=					[
								    [[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(_snake_case    ,			_snake_case)

								UpperCAmelCase_		:      List[Any]       	=					BatchSampler(range(22)    ,			batch_size=3    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      Dict       	=					[
								    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
								    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
								]
								self.check_batch_sampler_shards(_snake_case    ,			_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.
								UpperCAmelCase_		:      List[Any]       	=					BatchSampler(range(20)    ,			batch_size=3    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      Dict       	=					[
								    [[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(_snake_case    ,			_snake_case)

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

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

								UpperCAmelCase_		:      Any       	=					BatchSampler(range(2)    ,			batch_size=3    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      List[str]       	=					[[], []]
								self.check_batch_sampler_shards(_snake_case    ,			_snake_case)


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

								UpperCAmelCase_		:      Any       	=					BatchSampler(range(24)    ,			batch_size=4    ,			drop_last=_snake_case)
								# Expected shouldn't change
								self.check_batch_sampler_shards(_snake_case    ,			_snake_case    ,			split_batches=_snake_case)

								# Check the shards when the dataset is not a round multiple of batch size.
								UpperCAmelCase_		:      List[Any]       	=					BatchSampler(range(22)    ,			batch_size=4    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      Any       	=					[
								    [[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(_snake_case    ,			_snake_case    ,			split_batches=_snake_case)

								UpperCAmelCase_		:      int       	=					BatchSampler(range(22)    ,			batch_size=4    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      Any       	=					[
								    [[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(_snake_case    ,			_snake_case    ,			split_batches=_snake_case)

								# Check the shards when the dataset is not a round multiple of batch size or num_processes.
								UpperCAmelCase_		:      Dict       	=					BatchSampler(range(21)    ,			batch_size=4    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      Any       	=					[
								    [[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(_snake_case    ,			_snake_case    ,			split_batches=_snake_case)

								UpperCAmelCase_		:      Tuple       	=					BatchSampler(range(21)    ,			batch_size=4    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      Dict       	=					[
								    [[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(_snake_case    ,			_snake_case    ,			split_batches=_snake_case)

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

								UpperCAmelCase_		:      str       	=					BatchSampler(range(2)    ,			batch_size=4    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      Optional[int]       	=					[[], []]
								self.check_batch_sampler_shards(_snake_case    ,			_snake_case    ,			split_batches=_snake_case)


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

								UpperCAmelCase_		:      List[str]       	=					BatchSampler(range(24)    ,			batch_size=3    ,			drop_last=_snake_case)
								# Expected shouldn't change
								self.check_batch_sampler_shards(_snake_case    ,			_snake_case    ,			even_batches=_snake_case)

								# Check the shards when the dataset is a round multiple of batch size but not total batch size.
								UpperCAmelCase_		:      Optional[int]       	=					BatchSampler(range(21)    ,			batch_size=3    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      Union[str, Any]       	=					[
								    [[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(_snake_case    ,			_snake_case    ,			even_batches=_snake_case)

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

								# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
								# num_processes batch.
								UpperCAmelCase_		:      List[Any]       	=					BatchSampler(range(22)    ,			batch_size=3    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      str       	=					[
								    [[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(_snake_case    ,			_snake_case    ,			even_batches=_snake_case)

								UpperCAmelCase_		:      int       	=					BatchSampler(range(22)    ,			batch_size=3    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      List[str]       	=					[
								    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
								    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
								]
								self.check_batch_sampler_shards(_snake_case    ,			_snake_case    ,			even_batches=_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.
								UpperCAmelCase_		:      Union[str, Any]       	=					BatchSampler(range(20)    ,			batch_size=3    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      List[str]       	=					[
								    [[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(_snake_case    ,			_snake_case    ,			even_batches=_snake_case)

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

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

								UpperCAmelCase_		:      Union[str, Any]       	=					BatchSampler(range(2)    ,			batch_size=3    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      List[str]       	=					[[], []]
								self.check_batch_sampler_shards(_snake_case    ,			_snake_case    ,			even_batches=_snake_case)


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

								UpperCAmelCase_		:      List[str]       	=					BatchSampler(range(24)    ,			batch_size=4    ,			drop_last=_snake_case)
								# Expected shouldn't change
								self.check_batch_sampler_shards(_snake_case    ,			_snake_case    ,			split_batches=_snake_case    ,			even_batches=_snake_case)

								# Check the shards when the dataset is not a round multiple of batch size.
								UpperCAmelCase_		:      Optional[Any]       	=					BatchSampler(range(22)    ,			batch_size=4    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      Tuple       	=					[
								    [[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(_snake_case    ,			_snake_case    ,			split_batches=_snake_case    ,			even_batches=_snake_case)

								UpperCAmelCase_		:      str       	=					BatchSampler(range(22)    ,			batch_size=4    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      List[str]       	=					[
								    [[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(_snake_case    ,			_snake_case    ,			split_batches=_snake_case    ,			even_batches=_snake_case)

								# Check the shards when the dataset is not a round multiple of batch size or num_processes.
								UpperCAmelCase_		:      Optional[int]       	=					BatchSampler(range(21)    ,			batch_size=4    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      Union[str, Any]       	=					[
								    [[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(_snake_case    ,			_snake_case    ,			split_batches=_snake_case    ,			even_batches=_snake_case)

								UpperCAmelCase_		:      Optional[int]       	=					BatchSampler(range(21)    ,			batch_size=4    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      Dict       	=					[
								    [[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(_snake_case    ,			_snake_case    ,			split_batches=_snake_case    ,			even_batches=_snake_case)

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

								UpperCAmelCase_		:      Dict       	=					BatchSampler(range(2)    ,			batch_size=4    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      List[str]       	=					[[], []]
								self.check_batch_sampler_shards(_snake_case    ,			_snake_case    ,			split_batches=_snake_case    ,			even_batches=_snake_case)


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

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

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

								UpperCAmelCase_		:      Dict       	=					[]
								for idx in range(0    ,			len(_snake_case)    ,			_snake_case):
									for l in iterable_dataset_lists:
										observed += l[idx : idx + shard_batch_size]

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


							def 						_snake_case							(			self)							->				Tuple:
								UpperCAmelCase_		:      str       	=					42
								UpperCAmelCase_		:      Dict       	=					RandomIterableDataset()

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

								# Edge case with a very small dataset
								UpperCAmelCase_		:      Union[str, Any]       	=					RandomIterableDataset(max_length=2)

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


							def 						_snake_case							(			self)							->				Dict:
								UpperCAmelCase_		:      List[str]       	=					BatchSampler(range(16)    ,			batch_size=4    ,			drop_last=_snake_case)
								UpperCAmelCase_		:      Optional[int]       	=					SkipBatchSampler(_snake_case    ,			2)
								self.assertListEqual(list(_snake_case)    ,			[[8, 9, 10, 11], [12, 13, 14, 15]])


							def 						_snake_case							(			self)							->				Dict:
								UpperCAmelCase_		:      Tuple       	=					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)							->				Any:
								UpperCAmelCase_		:      Any       	=					DataLoader(list(range(16))    ,			batch_size=4)
								UpperCAmelCase_		:      List[str]       	=					skip_first_batches(_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)							->				List[Any]:
								UpperCAmelCase_		:      List[str]       	=					DataLoaderShard(list(range(16))    ,			batch_size=4)
								for idx, _ in enumerate(_snake_case):
									self.assertEqual(dataloader.end_of_dataloader    ,			idx == 3)

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





							def 						_snake_case							(			self)							->				Optional[int]:
								Accelerator()
								UpperCAmelCase_		:      List[Any]       	=					DataLoaderDispatcher(range(16)    ,			batch_size=4)
								for idx, _ in enumerate(_snake_case):
									self.assertEqual(dataloader.end_of_dataloader    ,			idx == 3)

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









'''simple docstring'''


import qiskit







def 						SCREAMING_SNAKE_CASE(  UpperCamelCase   ,UpperCamelCase  )    ->    qiskit.result.counts.Counts:
	UpperCAmelCase_		:      Tuple       	=					qiskit.Aer.get_backend('aer_simulator'  )

	UpperCAmelCase_		:      List[str]       	=					qiskit.QuantumCircuit(4   ,2  )
	# encode inputs in qubits 0 and 1
	if bita == 1:
		qc_ha.x(0  )
	if bita == 1:
		qc_ha.x(1  )
	qc_ha.barrier()

	# use cnots to write XOR of the inputs on qubit2
	qc_ha.cx(0   ,2  )
	qc_ha.cx(1   ,2  )

	# use ccx / toffoli gate to write AND of the inputs on qubit3
	qc_ha.ccx(0   ,1   ,3  )
	qc_ha.barrier()

	# extract outputs
	qc_ha.measure(2   ,0  )  # extract XOR value
	qc_ha.measure(3   ,1  )  # extract AND value

	# Execute the circuit on the qasm simulator
	UpperCAmelCase_		:      List[str]       	=					qiskit.execute(UpperCamelCase   ,UpperCamelCase   ,shots=1_0_0_0  )

	# Return the histogram data of the results of the experiment
	return job.result().get_counts(UpperCamelCase  )


if __name__ == "__main__":
						lowerCAmelCase__										= half_adder(1, 1)
						print(F'Half Adder Output Qubit Counts: {counts}')






