


'''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 		_snake_case					(     lowerCAmelCase__	):




					def __init__(      self							,							_lowerCamelCase=0.01							,							_lowerCamelCase=1000):
								UpperCAmelCase__  :					Optional[int]        =     p_stop
								UpperCAmelCase__  :					Union[str, Any]        =     max_length




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


class 		_snake_case					(     unittest.TestCase	):




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase=False							,							_lowerCamelCase=True):
								UpperCAmelCase__  :					Any        =     [
								    BatchSamplerShard(_lowerCamelCase							,							2							,							_lowerCamelCase							,							split_batches=_lowerCamelCase							,							even_batches=_lowerCamelCase)
								    for i in range(2)
								]
								UpperCAmelCase__  :					Union[str, Any]        =     [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 				snake_case__						(      self):
								UpperCAmelCase__  :					Any        =     BatchSampler(range(24)							,							batch_size=3							,							drop_last=_lowerCamelCase)
								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(_lowerCamelCase							,							_lowerCamelCase)

								UpperCAmelCase__  :					str        =     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__  :					List[Any]        =     BatchSampler(range(21)							,							batch_size=3							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					Dict        =     [
								    [[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__  :					Any        =     BatchSampler(range(21)							,							batch_size=3							,							drop_last=_lowerCamelCase)
								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(_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__  :					List[Any]        =     BatchSampler(range(22)							,							batch_size=3							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					Dict        =     [
								    [[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__  :					List[Any]        =     BatchSampler(range(22)							,							batch_size=3							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					Any        =     [
								    [[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__  :					Tuple        =     BatchSampler(range(20)							,							batch_size=3							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					Union[str, Any]        =     [
								    [[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__  :					List[str]        =     BatchSampler(range(20)							,							batch_size=3							,							drop_last=_lowerCamelCase)
								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(_lowerCamelCase							,							_lowerCamelCase)

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

								UpperCAmelCase__  :					int        =     BatchSampler(range(2)							,							batch_size=3							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					Optional[int]        =     [[], []]
								self.check_batch_sampler_shards(_lowerCamelCase							,							_lowerCamelCase)




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					List[str]        =     BatchSampler(range(24)							,							batch_size=4							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					int        =     [
								    [[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__  :					int        =     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__  :					List[str]        =     BatchSampler(range(22)							,							batch_size=4							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					List[str]        =     [
								    [[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__  :					Any        =     BatchSampler(range(22)							,							batch_size=4							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					int        =     [
								    [[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__  :					int        =     BatchSampler(range(21)							,							batch_size=4							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					Dict        =     [
								    [[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__  :					int        =     BatchSampler(range(21)							,							batch_size=4							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					int        =     [
								    [[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__  :					int        =     BatchSampler(range(2)							,							batch_size=4							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					Union[str, Any]        =     [[[0, 1]], [[0, 1]]]
								self.check_batch_sampler_shards(_lowerCamelCase							,							_lowerCamelCase							,							split_batches=_lowerCamelCase)

								UpperCAmelCase__  :					Dict        =     BatchSampler(range(2)							,							batch_size=4							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					Tuple        =     [[], []]
								self.check_batch_sampler_shards(_lowerCamelCase							,							_lowerCamelCase							,							split_batches=_lowerCamelCase)




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Union[str, Any]        =     BatchSampler(range(24)							,							batch_size=3							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					Tuple        =     [
								    [[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__  :					Optional[Any]        =     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__  :					List[Any]        =     BatchSampler(range(21)							,							batch_size=3							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					Optional[int]        =     [
								    [[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__  :					List[str]        =     BatchSampler(range(21)							,							batch_size=3							,							drop_last=_lowerCamelCase)
								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(_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__  :					Optional[Any]        =     BatchSampler(range(22)							,							batch_size=3							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					List[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(_lowerCamelCase							,							_lowerCamelCase							,							even_batches=_lowerCamelCase)

								UpperCAmelCase__  :					Union[str, Any]        =     BatchSampler(range(22)							,							batch_size=3							,							drop_last=_lowerCamelCase)
								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(_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__  :					List[str]        =     BatchSampler(range(20)							,							batch_size=3							,							drop_last=_lowerCamelCase)
								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(_lowerCamelCase							,							_lowerCamelCase							,							even_batches=_lowerCamelCase)

								UpperCAmelCase__  :					List[str]        =     BatchSampler(range(20)							,							batch_size=3							,							drop_last=_lowerCamelCase)
								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(_lowerCamelCase							,							_lowerCamelCase							,							even_batches=_lowerCamelCase)

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

								UpperCAmelCase__  :					int        =     BatchSampler(range(2)							,							batch_size=3							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					str        =     [[], []]
								self.check_batch_sampler_shards(_lowerCamelCase							,							_lowerCamelCase							,							even_batches=_lowerCamelCase)




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					List[str]        =     BatchSampler(range(24)							,							batch_size=4							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					List[str]        =     [
								    [[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__  :					str        =     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__  :					Union[str, Any]        =     BatchSampler(range(22)							,							batch_size=4							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					Dict        =     [
								    [[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__  :					Any        =     BatchSampler(range(22)							,							batch_size=4							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					List[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(_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__  :					Dict        =     BatchSampler(range(21)							,							batch_size=4							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					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(_lowerCamelCase							,							_lowerCamelCase							,							split_batches=_lowerCamelCase							,							even_batches=_lowerCamelCase)

								UpperCAmelCase__  :					Any        =     BatchSampler(range(21)							,							batch_size=4							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					int        =     [
								    [[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__  :					Any        =     BatchSampler(range(2)							,							batch_size=4							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					str        =     [[[0, 1]], []]
								self.check_batch_sampler_shards(_lowerCamelCase							,							_lowerCamelCase							,							split_batches=_lowerCamelCase							,							even_batches=_lowerCamelCase)

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




					def 				snake_case__						(      self):
								UpperCAmelCase__  :					Dict        =     [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
								UpperCAmelCase__  :					Tuple        =     [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 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase=False							,							_lowerCamelCase=2							,							_lowerCamelCase=False):
								random.seed(_lowerCamelCase)
								UpperCAmelCase__  :					Dict        =     list(_lowerCamelCase)

								UpperCAmelCase__  :					Union[str, Any]        =     [
								    IterableDatasetShard(
								        _lowerCamelCase							,							batch_size=_lowerCamelCase							,							drop_last=_lowerCamelCase							,							num_processes=_lowerCamelCase							,							process_index=_lowerCamelCase							,							split_batches=_lowerCamelCase							,							)
								    for i in range(_lowerCamelCase)
								]
								UpperCAmelCase__  :					Any        =     []
								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__  :					Optional[int]        =     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__  :					str        =     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__  :					Tuple        =     []
								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 				snake_case__						(      self):
								UpperCAmelCase__  :					Union[str, Any]        =     42
								UpperCAmelCase__  :					List[Any]        =     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__  :					Any        =     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 				snake_case__						(      self):
								UpperCAmelCase__  :					Tuple        =     BatchSampler(range(16)							,							batch_size=4							,							drop_last=_lowerCamelCase)
								UpperCAmelCase__  :					Dict        =     SkipBatchSampler(_lowerCamelCase							,							2)
								self.assertListEqual(list(_lowerCamelCase)							,							[[8, 9, 10, 11], [12, 13, 14, 15]])




					def 				snake_case__						(      self):
								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):
								UpperCAmelCase__  :					str        =     DataLoader(list(range(16))							,							batch_size=4)
								UpperCAmelCase__  :					Optional[Any]        =     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 				snake_case__						(      self):
								UpperCAmelCase__  :					Tuple        =     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 				snake_case__						(      self):
								Accelerator()
								UpperCAmelCase__  :					Union[str, Any]        =     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)



'''simple docstring'''



import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
__A	=object()

# For specifying empty leaf dict `{}`
__A	=object()







def     _UpperCamelCase      (		UpperCamelCase__ ,      UpperCamelCase__					):
			UpperCAmelCase__  :					Union[str, Any]        =     tuple((re.compile(x + """$"""					) for x in qs)					)
			for i in range(len(UpperCamelCase__					) - len(UpperCamelCase__					) + 1					):
						UpperCAmelCase__  :					Optional[Any]        =     [x.match(UpperCamelCase__					) for x, y in zip(UpperCamelCase__ ,      ks[i:]					)]
						if matches and all(UpperCamelCase__					):
									return True
			return False







def     _UpperCamelCase      (		UpperCamelCase__					):
			def replace(UpperCamelCase__ ,      UpperCamelCase__					):
						for rule, replacement in rules:
									if _match(UpperCamelCase__ ,      UpperCamelCase__					):
												return replacement
						return val

			return replace







def     _UpperCamelCase      (		):
			return [
			    # embeddings
			    (("transformer", "wpe", "embedding"), P("""mp""" ,      UpperCamelCase__					)),
			    (("transformer", "wte", "embedding"), P("""mp""" ,      UpperCamelCase__					)),
			    # atention
			    (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(UpperCamelCase__ ,      """mp"""					)),
			    (("attention", "out_proj", "kernel"), P("""mp""" ,      UpperCamelCase__					)),
			    (("attention", "out_proj", "bias"), None),
			    # mlp
			    (("mlp", "c_fc", "kernel"), P(UpperCamelCase__ ,      """mp"""					)),
			    (("mlp", "c_fc", "bias"), P("""mp"""					)),
			    (("mlp", "c_proj", "kernel"), P("""mp""" ,      UpperCamelCase__					)),
			    (("mlp", "c_proj", "bias"), None),
			    # layer norms
			    ((r"ln_\d+", "bias"), None),
			    ((r"\d+", r"ln_\d+", "scale"), None),
			    (("ln_f", "bias"), None),
			    (("ln_f", "scale"), None),
			]







def     _UpperCamelCase      (		UpperCamelCase__					):
			UpperCAmelCase__  :					Optional[int]        =     _get_partition_rules()
			UpperCAmelCase__  :					Optional[int]        =     _replacement_rules(UpperCamelCase__					)
			UpperCAmelCase__  :					List[Any]        =     {k: _unmatched for k in flatten_dict(UpperCamelCase__					)}
			UpperCAmelCase__  :					Tuple        =     {k: replace(UpperCamelCase__ ,      UpperCamelCase__					) for k, v in initd.items()}
			assert _unmatched not in result.values(), "Incomplete partition spec."
			return freeze(unflatten_dict(UpperCamelCase__					)					)