

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




	def __init__(		self	:						Optional[Any]		,	lowerCAmelCase	:						List[Any]=0.01		,	lowerCAmelCase	:						Tuple=1_000							):


		'''simple docstring'''







		UpperCAmelCase_       				=			p_stop
		UpperCAmelCase_       				=			max_length




	def __iter__(		self	:						List[Any]							):


		'''simple docstring'''







		UpperCAmelCase_       				=			0
		UpperCAmelCase_       				=			False
		while not stop and count < self.max_length:
			yield count
			count += 1
			UpperCAmelCase_       				=			random.random() < self.p_stop


class 						__UpperCamelCase    (					unittest.TestCase						):




	def       __A					(		self	:						Tuple		,	lowerCAmelCase	:						Optional[Any]		,	lowerCAmelCase	:						List[str]		,	lowerCAmelCase	:						List[str]=False		,	lowerCAmelCase	:						Any=True							):


		'''simple docstring'''







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




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


		'''simple docstring'''







		UpperCAmelCase_       				=			BatchSampler(range(24							)		,	batch_size=3		,	drop_last=lowerCAmelCase							)
		UpperCAmelCase_       				=			[
		    [[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(lowerCAmelCase		,	lowerCAmelCase							)

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

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

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

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

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

		# 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_       				=			BatchSampler(range(20							)		,	batch_size=3		,	drop_last=lowerCAmelCase							)
		UpperCAmelCase_       				=			[
		    [[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(lowerCAmelCase		,	lowerCAmelCase							)

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

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

		UpperCAmelCase_       				=			BatchSampler(range(2							)		,	batch_size=3		,	drop_last=lowerCAmelCase							)
		UpperCAmelCase_       				=			[[], []]
		self.check_batch_sampler_shards(lowerCAmelCase		,	lowerCAmelCase							)




	def       __A					(		self	:						Dict							):


		'''simple docstring'''







		UpperCAmelCase_       				=			BatchSampler(range(24							)		,	batch_size=4		,	drop_last=lowerCAmelCase							)
		UpperCAmelCase_       				=			[
		    [[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(lowerCAmelCase		,	lowerCAmelCase		,	split_batches=lowerCAmelCase							)

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

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

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

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

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

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

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




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


		'''simple docstring'''







		UpperCAmelCase_       				=			BatchSampler(range(24							)		,	batch_size=3		,	drop_last=lowerCAmelCase							)
		UpperCAmelCase_       				=			[
		    [[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(lowerCAmelCase		,	lowerCAmelCase		,	even_batches=lowerCAmelCase							)

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

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

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

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

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

		# 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_       				=			BatchSampler(range(20							)		,	batch_size=3		,	drop_last=lowerCAmelCase							)
		UpperCAmelCase_       				=			[
		    [[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(lowerCAmelCase		,	lowerCAmelCase		,	even_batches=lowerCAmelCase							)

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

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

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




	def       __A					(		self	:						Dict							):


		'''simple docstring'''







		UpperCAmelCase_       				=			BatchSampler(range(24							)		,	batch_size=4		,	drop_last=lowerCAmelCase							)
		UpperCAmelCase_       				=			[
		    [[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(lowerCAmelCase		,	lowerCAmelCase		,	split_batches=lowerCAmelCase		,	even_batches=lowerCAmelCase							)

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

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

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

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

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

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

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




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


		'''simple docstring'''







		UpperCAmelCase_       				=			[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
		UpperCAmelCase_       				=			[BatchSamplerShard(lowerCAmelCase		,	2		,	lowerCAmelCase		,	even_batches=lowerCAmelCase							) 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	:						Union[str, Any]		,	lowerCAmelCase	:						Any		,	lowerCAmelCase	:						Tuple		,	lowerCAmelCase	:						Optional[int]		,	lowerCAmelCase	:						str=False		,	lowerCAmelCase	:						Union[str, Any]=2		,	lowerCAmelCase	:						Optional[int]=False							):


		'''simple docstring'''







		random.seed(lowerCAmelCase							)
		UpperCAmelCase_       				=			list(lowerCAmelCase							)

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

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

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

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




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


		'''simple docstring'''







		UpperCAmelCase_       				=			42
		UpperCAmelCase_       				=			RandomIterableDataset()

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

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

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




	def       __A					(		self	:						str							):


		'''simple docstring'''







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




	def       __A					(		self	:						int							):


		'''simple docstring'''







		UpperCAmelCase_       				=			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	:						Dict							):


		'''simple docstring'''







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




	def       __A					(		self	:						Any							):


		'''simple docstring'''







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

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




	def       __A					(		self	:						str							):


		'''simple docstring'''







		Accelerator()
		UpperCAmelCase_       				=			DataLoaderDispatcher(range(16							)		,	batch_size=4							)
		for idx, _ in enumerate(lowerCAmelCase							):
			self.assertEqual(dataloader.end_of_dataloader		,	idx == 3							)

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


from .imports import is_tqdm_available


if is_tqdm_available():
				from tqdm.auto import tqdm as _tqdm

from ..state import PartialState



def       __lowerCAmelCase      (						A = True       ,     *A       ,     **A    ):
	if not is_tqdm_available():
		raise ImportError("Accelerate's `tqdm` module requires `tqdm` to be installed. Please run `pip install tqdm`."    )
	UpperCAmelCase_       				=			False
	if main_process_only:
		UpperCAmelCase_       				=			PartialState().local_process_index == 0
	return _tqdm(*A       ,     **A       ,     disable=A    )