





import random
import unittest

from torch.utils.data import BatchSampler, DataLoader, IterableDataset

from accelerate import Accelerator
from accelerate.data_loader import (
    BatchSamplerShard,
    DataLoaderDispatcher,
    DataLoaderShard,
    IterableDatasetShard,
    SkipBatchSampler,
    SkipDataLoader,
    skip_first_batches,
)



class        _lowerCAmelCase    (			UpperCAmelCase_							):


  '''simple docstring'''







  def __init__(      self      :					List[Any]      ,				UpperCamelCase      :					Union[str, Any]=0.01      ,				UpperCamelCase      :					Optional[Any]=10_00							):


    '''simple docstring'''





    _snake_case   :				Tuple														=	p_stop
    _snake_case   :				Dict														=	max_length







  def __iter__(      self      :					Any							):


    '''simple docstring'''





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




class        _lowerCAmelCase    (			unittest.TestCase							):


  '''simple docstring'''







  def UpperCamelCase_ (      self      :					Optional[Any]      ,				UpperCamelCase      :					Dict      ,				UpperCamelCase      :					Dict      ,				UpperCamelCase      :					Optional[Any]=False      ,				UpperCamelCase      :					List[str]=True							):


    '''simple docstring'''





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







  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    _snake_case   :				Optional[Any]														=	BatchSampler(range(24							)      ,				batch_size=3      ,				drop_last=UpperCamelCase							)
    _snake_case   :				List[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(UpperCamelCase      ,				UpperCamelCase							)

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

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

    _snake_case   :				Any														=	BatchSampler(range(21							)      ,				batch_size=3      ,				drop_last=UpperCamelCase							)
    _snake_case   :				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(UpperCamelCase      ,				UpperCamelCase							)

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

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

    # 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   :				Union[str, Any]														=	BatchSampler(range(20							)      ,				batch_size=3      ,				drop_last=UpperCamelCase							)
    _snake_case   :				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(UpperCamelCase      ,				UpperCamelCase							)

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

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

    _snake_case   :				Dict														=	BatchSampler(range(2							)      ,				batch_size=3      ,				drop_last=UpperCamelCase							)
    _snake_case   :				Optional[int]														=	[[], []]
    self.check_batch_sampler_shards(UpperCamelCase      ,				UpperCamelCase							)







  def UpperCamelCase_ (      self      :					Tuple							):


    '''simple docstring'''





    _snake_case   :				Optional[Any]														=	BatchSampler(range(24							)      ,				batch_size=4      ,				drop_last=UpperCamelCase							)
    _snake_case   :				Optional[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(UpperCamelCase      ,				UpperCamelCase      ,				split_batches=UpperCamelCase							)

    _snake_case   :				Optional[int]														=	BatchSampler(range(24							)      ,				batch_size=4      ,				drop_last=UpperCamelCase							)
    # Expected shouldn't change
    self.check_batch_sampler_shards(UpperCamelCase      ,				UpperCamelCase      ,				split_batches=UpperCamelCase							)

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

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

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

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

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

    _snake_case   :				Dict														=	BatchSampler(range(2							)      ,				batch_size=4      ,				drop_last=UpperCamelCase							)
    _snake_case   :				Dict														=	[[], []]
    self.check_batch_sampler_shards(UpperCamelCase      ,				UpperCamelCase      ,				split_batches=UpperCamelCase							)







  def UpperCamelCase_ (      self      :					Tuple							):


    '''simple docstring'''





    _snake_case   :				Optional[Any]														=	BatchSampler(range(24							)      ,				batch_size=3      ,				drop_last=UpperCamelCase							)
    _snake_case   :				Dict														=	[
        [[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(UpperCamelCase      ,				UpperCamelCase      ,				even_batches=UpperCamelCase							)

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

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

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

    # 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(22							)      ,				batch_size=3      ,				drop_last=UpperCamelCase							)
    _snake_case   :				Optional[int]														=	[
        [[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(UpperCamelCase      ,				UpperCamelCase      ,				even_batches=UpperCamelCase							)

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

    # 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   :				Union[str, Any]														=	BatchSampler(range(20							)      ,				batch_size=3      ,				drop_last=UpperCamelCase							)
    _snake_case   :				Tuple														=	[
        [[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(UpperCamelCase      ,				UpperCamelCase      ,				even_batches=UpperCamelCase							)

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

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

    _snake_case   :				Any														=	BatchSampler(range(2							)      ,				batch_size=3      ,				drop_last=UpperCamelCase							)
    _snake_case   :				int														=	[[], []]
    self.check_batch_sampler_shards(UpperCamelCase      ,				UpperCamelCase      ,				even_batches=UpperCamelCase							)







  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    _snake_case   :				Dict														=	BatchSampler(range(24							)      ,				batch_size=4      ,				drop_last=UpperCamelCase							)
    _snake_case   :				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(UpperCamelCase      ,				UpperCamelCase      ,				split_batches=UpperCamelCase      ,				even_batches=UpperCamelCase							)

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

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

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

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

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

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

    _snake_case   :				Any														=	BatchSampler(range(2							)      ,				batch_size=4      ,				drop_last=UpperCamelCase							)
    _snake_case   :				Union[str, Any]														=	[[], []]
    self.check_batch_sampler_shards(UpperCamelCase      ,				UpperCamelCase      ,				split_batches=UpperCamelCase      ,				even_batches=UpperCamelCase							)







  def UpperCamelCase_ (      self      :					List[str]							):


    '''simple docstring'''





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


    '''simple docstring'''





    random.seed(UpperCamelCase							)
    _snake_case   :				Optional[Any]														=	list(UpperCamelCase							)

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

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

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

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







  def UpperCamelCase_ (      self      :					List[str]							):


    '''simple docstring'''





    _snake_case   :				List[str]														=	42
    _snake_case   :				List[Any]														=	RandomIterableDataset()

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

    # Edge case with a very small dataset
    _snake_case   :				Optional[Any]														=	RandomIterableDataset(max_length=2							)

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







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


    '''simple docstring'''





    _snake_case   :				Optional[int]														=	BatchSampler(range(16							)      ,				batch_size=4      ,				drop_last=UpperCamelCase							)
    _snake_case   :				int														=	SkipBatchSampler(UpperCamelCase      ,				2							)
    self.assertListEqual(list(UpperCamelCase							)      ,				[[8, 9, 10, 11], [12, 13, 14, 15]]							)







  def UpperCamelCase_ (      self      :					Dict							):


    '''simple docstring'''





    _snake_case   :				Optional[int]														=	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 UpperCamelCase_ (      self      :					Optional[int]							):


    '''simple docstring'''





    _snake_case   :				Optional[Any]														=	DataLoader(list(range(16							)							)      ,				batch_size=4							)
    _snake_case   :				int														=	skip_first_batches(UpperCamelCase      ,				num_batches=2							)
    self.assertListEqual([t.tolist() for t in new_dataloader]      ,				[[8, 9, 10, 11], [12, 13, 14, 15]]							)







  def UpperCamelCase_ (      self      :					Tuple							):


    '''simple docstring'''





    _snake_case   :				Optional[int]														=	DataLoaderShard(list(range(16							)							)      ,				batch_size=4							)
    for idx, _ in enumerate(UpperCamelCase							):
      self.assertEqual(dataloader.end_of_dataloader      ,				idx == 3							)

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







  def UpperCamelCase_ (      self      :					Any							):


    '''simple docstring'''





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

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











from unittest import TestCase

from datasets import Dataset
from minhash_deduplication import deduplicate_dataset, make_duplicate_clusters


def     lowerCamelCase_	(      )->		Any:
  _snake_case   :				List[str]														=	{
      'repo_name': ['test_repo1', 'test_repo2', 'test_repo3'],
      'path': ['test_1.py', 'test_2.py', 'unit_test.py'],
      'content': ['a ' * 20, 'a ' * 30, 'b ' * 7],
  }
  _snake_case   :				Optional[Any]														=	Dataset.from_dict(lowerCAmelCase				)
  return dataset







class        _lowerCAmelCase    (			UpperCAmelCase_							):


  '''simple docstring'''







  def UpperCamelCase_ (      self      :					Dict							):


    '''simple docstring'''





    _snake_case   :				Union[str, Any]														=	get_dataset()
    _snake_case   :				Tuple														=	make_duplicate_clusters(UpperCamelCase      ,				0.85							)
    self.assertEqual(len(duplicate_clusters[0]							)      ,				2							)







  def UpperCamelCase_ (      self      :					Dict							):


    '''simple docstring'''





    _snake_case   :				List[str]														=	get_dataset()
    _snake_case							,					_snake_case   :				str														=	deduplicate_dataset(UpperCamelCase							)
    self.assertEqual(len(UpperCamelCase							)      ,				2							)
    print(UpperCamelCase							)
    self.assertEqual(duplicate_clusters[0][0]['copies']      ,				2							)
    self.assertEqual(duplicate_clusters[0][0]['is_extreme']      ,				UpperCamelCase							)




