
from collections import OrderedDict
from typing import Any, List, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import logging


__lowerCAmelCase     :	List[str]       			=  logging.get_logger(__name__)


__lowerCAmelCase     :	str       			=  {
    'Salesforce/codegen-350M-nl': 'https://huggingface.co/Salesforce/codegen-350M-nl/resolve/main/config.json',
    'Salesforce/codegen-350M-multi': 'https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/config.json',
    'Salesforce/codegen-350M-mono': 'https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/config.json',
    'Salesforce/codegen-2B-nl': 'https://huggingface.co/Salesforce/codegen-2B-nl/resolve/main/config.json',
    'Salesforce/codegen-2B-multi': 'https://huggingface.co/Salesforce/codegen-2B-multi/resolve/main/config.json',
    'Salesforce/codegen-2B-mono': 'https://huggingface.co/Salesforce/codegen-2B-mono/resolve/main/config.json',
    'Salesforce/codegen-6B-nl': 'https://huggingface.co/Salesforce/codegen-6B-nl/resolve/main/config.json',
    'Salesforce/codegen-6B-multi': 'https://huggingface.co/Salesforce/codegen-6B-multi/resolve/main/config.json',
    'Salesforce/codegen-6B-mono': 'https://huggingface.co/Salesforce/codegen-6B-mono/resolve/main/config.json',
    'Salesforce/codegen-16B-nl': 'https://huggingface.co/Salesforce/codegen-16B-nl/resolve/main/config.json',
    'Salesforce/codegen-16B-multi': 'https://huggingface.co/Salesforce/codegen-16B-multi/resolve/main/config.json',
    'Salesforce/codegen-16B-mono': 'https://huggingface.co/Salesforce/codegen-16B-mono/resolve/main/config.json',
}




class 			UpperCAmelCase_     (					_A	):


  '''simple docstring'''







  a__       			=						"""codegen"""
  a__       			=						{
      """max_position_embeddings""": """n_positions""",
      """hidden_size""": """n_embd""",
      """num_attention_heads""": """n_head""",
      """num_hidden_layers""": """n_layer""",
  }
  def __init__(					self     :					List[str]     ,				UpperCamelCase__     :					Optional[Any]=5_0400     ,				UpperCamelCase__     :					str=2048     ,				UpperCamelCase__     :					Union[str, Any]=2048     ,				UpperCamelCase__     :					Optional[int]=4096     ,				UpperCamelCase__     :					Any=28     ,				UpperCamelCase__     :					Any=16     ,				UpperCamelCase__     :					Any=64     ,				UpperCamelCase__     :					Optional[int]=None     ,				UpperCamelCase__     :					Optional[int]="gelu_new"     ,				UpperCamelCase__     :					Optional[Any]=0.0     ,				UpperCamelCase__     :					Dict=0.0     ,				UpperCamelCase__     :					Union[str, Any]=0.0     ,				UpperCamelCase__     :					Dict=1E-5     ,				UpperCamelCase__     :					Tuple=0.02     ,				UpperCamelCase__     :					Optional[int]=True     ,				UpperCamelCase__     :					List[str]=5_0256     ,				UpperCamelCase__     :					Tuple=5_0256     ,				UpperCamelCase__     :					str=False     ,				**UpperCamelCase__     :					Dict     ,				)     ->      Union[str, Any]:







    """simple docstring"""






    __magic_name__								=			vocab_size
    __magic_name__								=			n_ctx
    __magic_name__								=			n_positions
    __magic_name__								=			n_embd
    __magic_name__								=			n_layer
    __magic_name__								=			n_head
    __magic_name__								=			n_inner
    __magic_name__								=			rotary_dim
    __magic_name__								=			activation_function
    __magic_name__								=			resid_pdrop
    __magic_name__								=			embd_pdrop
    __magic_name__								=			attn_pdrop
    __magic_name__								=			layer_norm_epsilon
    __magic_name__								=			initializer_range
    __magic_name__								=			use_cache

    __magic_name__								=			bos_token_id
    __magic_name__								=			eos_token_id

    super().__init__(
        bos_token_id=UpperCamelCase__     ,				eos_token_id=UpperCamelCase__     ,				tie_word_embeddings=UpperCamelCase__     ,				**UpperCamelCase__       )







class 			UpperCAmelCase_     (					_A	):


  '''simple docstring'''







  def __init__(					self     :					int     ,				UpperCamelCase__     :					PretrainedConfig     ,				UpperCamelCase__     :					str = "default"     ,				UpperCamelCase__     :					List[PatchingSpec] = None     ,				UpperCamelCase__     :					bool = False     ,				)     ->      Optional[int]:







    """simple docstring"""






    super().__init__(UpperCamelCase__     ,				task=UpperCamelCase__     ,				patching_specs=UpperCamelCase__     ,				use_past=UpperCamelCase__       )
    if not getattr(self._config     ,				"""pad_token_id"""     ,				UpperCamelCase__       ):
      # TODO: how to do that better?
      __magic_name__								=			0
  @property
  def 				_lowercase				(					self     :					Any       )     ->      Mapping[str, Mapping[int, str]]:







    """simple docstring"""






    __magic_name__								=			OrderedDict({"""input_ids""": {0: """batch""", 1: """sequence"""}}       )
    if self.use_past:
      self.fill_with_past_key_values_(UpperCamelCase__     ,				direction="""inputs"""       )
      __magic_name__								=			{0: """batch""", 1: """past_sequence + sequence"""}
    else:
      __magic_name__								=			{0: """batch""", 1: """sequence"""}

    return common_inputs
  @property
  def 				_lowercase				(					self     :					Optional[Any]       )     ->      int:







    """simple docstring"""






    return self._config.n_layer
  @property
  def 				_lowercase				(					self     :					List[Any]       )     ->      int:







    """simple docstring"""






    return self._config.n_head
  def 				_lowercase				(					self     :					List[str]     ,				UpperCamelCase__     :					PreTrainedTokenizer     ,				UpperCamelCase__     :					int = -1     ,				UpperCamelCase__     :					int = -1     ,				UpperCamelCase__     :					bool = False     ,				UpperCamelCase__     :					Optional[TensorType] = None     ,				)     ->      Mapping[str, Any]:







    """simple docstring"""






    __magic_name__								=			super(UpperCamelCase__     ,				self       ).generate_dummy_inputs(
        UpperCamelCase__     ,				batch_size=UpperCamelCase__     ,				seq_length=UpperCamelCase__     ,				is_pair=UpperCamelCase__     ,				framework=UpperCamelCase__       )

    # We need to order the input in the way they appears in the forward()
    __magic_name__								=			OrderedDict({"""input_ids""": common_inputs["""input_ids"""]}       )

    # Need to add the past_keys
    if self.use_past:
      if not is_torch_available():
        raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""       )
      else:
        import torch

        __magic_name__	,       __magic_name__								=			common_inputs["""input_ids"""].shape
        # Not using the same length for past_key_values
        __magic_name__								=			seqlen + 2
        __magic_name__								=			(
            batch,
            self.num_attention_heads,
            past_key_values_length,
            self._config.hidden_size // self.num_attention_heads,
        )
        __magic_name__								=			[
            (torch.zeros(UpperCamelCase__       ), torch.zeros(UpperCamelCase__       )) for _ in range(self.num_layers       )
        ]

    __magic_name__								=			common_inputs["""attention_mask"""]
    if self.use_past:
      __magic_name__								=			ordered_inputs["""attention_mask"""].dtype
      __magic_name__								=			torch.cat(
          [ordered_inputs["""attention_mask"""], torch.ones(UpperCamelCase__     ,				UpperCamelCase__     ,				dtype=UpperCamelCase__       )]     ,				dim=1       )

    return ordered_inputs





  @property
  def 				_lowercase				(					self     :					Optional[int]       )     ->      int:







    """simple docstring"""






    return 13



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_     (					_A	):


  '''simple docstring'''







  def __init__(					self     :					Optional[int]     ,				UpperCamelCase__     :					Optional[Any]=0.01     ,				UpperCamelCase__     :					Tuple=1000       )     ->      str:







    """simple docstring"""






    __magic_name__								=			p_stop
    __magic_name__								=			max_length





  def __iter__(					self     :					Union[str, Any]       )     ->      Dict:







    """simple docstring"""






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







class 			UpperCAmelCase_     (					unittest.TestCase	):


  '''simple docstring'''







  def 				_lowercase				(					self     :					Optional[Any]     ,				UpperCamelCase__     :					Union[str, Any]     ,				UpperCamelCase__     :					Tuple     ,				UpperCamelCase__     :					str=False     ,				UpperCamelCase__     :					Optional[Any]=True       )     ->      Dict:







    """simple docstring"""






    __magic_name__								=			[
        BatchSamplerShard(UpperCamelCase__     ,				2     ,				UpperCamelCase__     ,				split_batches=UpperCamelCase__     ,				even_batches=UpperCamelCase__       )
        for i in range(2       )
    ]
    __magic_name__								=			[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 				_lowercase				(					self     :					Any       )     ->      Union[str, Any]:







    """simple docstring"""






    __magic_name__								=			BatchSampler(range(24       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			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.
    __magic_name__								=			BatchSampler(range(21       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			BatchSampler(range(21       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(22       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			BatchSampler(range(22       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(20       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			BatchSampler(range(20       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[[0, 1, 0]], [[1, 0, 1]]]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__       )

    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[], []]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__       )
  def 				_lowercase				(					self     :					Tuple       )     ->      int:







    """simple docstring"""






    __magic_name__								=			BatchSampler(range(24       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			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.
    __magic_name__								=			BatchSampler(range(22       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			BatchSampler(range(22       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(21       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			BatchSampler(range(21       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[[0, 1]], [[0, 1]]]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				split_batches=UpperCamelCase__       )

    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[], []]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				split_batches=UpperCamelCase__       )
  def 				_lowercase				(					self     :					int       )     ->      str:







    """simple docstring"""






    __magic_name__								=			BatchSampler(range(24       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			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.
    __magic_name__								=			BatchSampler(range(21       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			BatchSampler(range(21       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(22       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			BatchSampler(range(22       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(20       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			BatchSampler(range(20       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[[0, 1]], []]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				even_batches=UpperCamelCase__       )

    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[], []]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				even_batches=UpperCamelCase__       )
  def 				_lowercase				(					self     :					Optional[int]       )     ->      int:







    """simple docstring"""






    __magic_name__								=			BatchSampler(range(24       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			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.
    __magic_name__								=			BatchSampler(range(22       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			BatchSampler(range(22       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(21       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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__       )

    __magic_name__								=			BatchSampler(range(21       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[
        [[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.
    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[[0, 1]], []]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				split_batches=UpperCamelCase__     ,				even_batches=UpperCamelCase__       )

    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[], []]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				split_batches=UpperCamelCase__     ,				even_batches=UpperCamelCase__       )
  def 				_lowercase				(					self     :					Dict       )     ->      str:







    """simple docstring"""






    __magic_name__								=			[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
    __magic_name__								=			[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 				_lowercase				(					self     :					int     ,				UpperCamelCase__     :					List[str]     ,				UpperCamelCase__     :					Union[str, Any]     ,				UpperCamelCase__     :					Tuple     ,				UpperCamelCase__     :					Dict=False     ,				UpperCamelCase__     :					Union[str, Any]=2     ,				UpperCamelCase__     :					Union[str, Any]=False       )     ->      Union[str, Any]:







    """simple docstring"""






    random.seed(UpperCamelCase__       )
    __magic_name__								=			list(UpperCamelCase__       )

    __magic_name__								=			[
        IterableDatasetShard(
            UpperCamelCase__     ,				batch_size=UpperCamelCase__     ,				drop_last=UpperCamelCase__     ,				num_processes=UpperCamelCase__     ,				process_index=UpperCamelCase__     ,				split_batches=UpperCamelCase__     ,				)
        for i in range(UpperCamelCase__       )
    ]
    __magic_name__								=			[]
    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__       )       )

    __magic_name__								=			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
    __magic_name__								=			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       )

    __magic_name__								=			[]
    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 				_lowercase				(					self     :					Tuple       )     ->      List[str]:







    """simple docstring"""






    __magic_name__								=			42
    __magic_name__								=			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
    __magic_name__								=			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 				_lowercase				(					self     :					List[Any]       )     ->      int:







    """simple docstring"""






    __magic_name__								=			BatchSampler(range(16       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			SkipBatchSampler(UpperCamelCase__     ,				2       )
    self.assertListEqual(list(UpperCamelCase__       )     ,				[[8, 9, 10, 11], [12, 13, 14, 15]]       )
  def 				_lowercase				(					self     :					Tuple       )     ->      str:







    """simple docstring"""






    __magic_name__								=			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 				_lowercase				(					self     :					str       )     ->      List[str]:







    """simple docstring"""






    __magic_name__								=			DataLoader(list(range(16       )       )     ,				batch_size=4       )
    __magic_name__								=			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 				_lowercase				(					self     :					List[str]       )     ->      Tuple:







    """simple docstring"""






    __magic_name__								=			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 				_lowercase				(					self     :					Tuple       )     ->      List[str]:







    """simple docstring"""






    Accelerator()
    __magic_name__								=			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       )

