

"""simple docstring"""





from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import tensorflow as tf

   from transformers.models.distilbert.modeling_tf_distilbert import (
       TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
       TFDistilBertForMaskedLM,
       TFDistilBertForMultipleChoice,
       TFDistilBertForQuestionAnswering,
       TFDistilBertForSequenceClassification,
       TFDistilBertForTokenClassification,
       TFDistilBertModel,
   )




class 					_lowerCAmelCase  :







 """simple docstring"""







 def __init__(			self							:			List[str],					UpperCAmelCase__							:			Tuple,					):
    __lowercase								=						parent
    __lowercase								=						1_3
    __lowercase								=						7
    __lowercase								=						True
    __lowercase								=						True
    __lowercase								=						False
    __lowercase								=						True
    __lowercase								=						9_9
    __lowercase								=						3_2
    __lowercase								=						2
    __lowercase								=						4
    __lowercase								=						3_7
    __lowercase								=						"gelu"
    __lowercase								=						0.1
    __lowercase								=						0.1
    __lowercase								=						5_1_2
    __lowercase								=						1_6
    __lowercase								=						2
    __lowercase								=						0.02
    __lowercase								=						3
    __lowercase								=						4
    __lowercase								=						None


 def 					_lowercase (			self							:			int	):
    __lowercase								=						ids_tensor([self.batch_size, self.seq_length],					self.vocab_size	)

    __lowercase								=						None
    if self.use_input_mask:
       __lowercase								=						random_attention_mask([self.batch_size, self.seq_length]	)

    __lowercase								=						None
    __lowercase								=						None
    __lowercase								=						None
    if self.use_labels:
       __lowercase								=						ids_tensor([self.batch_size],					self.type_sequence_label_size	)
       __lowercase								=						ids_tensor([self.batch_size, self.seq_length],					self.num_labels	)
       __lowercase								=						ids_tensor([self.batch_size],					self.num_choices	)

    __lowercase								=						DistilBertConfig(
        vocab_size=self.vocab_size,					dim=self.hidden_size,					n_layers=self.num_hidden_layers,					n_heads=self.num_attention_heads,					hidden_dim=self.intermediate_size,					hidden_act=self.hidden_act,					dropout=self.hidden_dropout_prob,					attention_dropout=self.attention_probs_dropout_prob,					max_position_embeddings=self.max_position_embeddings,					initializer_range=self.initializer_range,					)

    return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels


 def 					_lowercase (			self							:			Union[str, Any],					UpperCAmelCase__							:			str,					UpperCAmelCase__							:			Dict,					UpperCAmelCase__							:			int,					UpperCAmelCase__							:			Dict,					UpperCAmelCase__							:			Any,					UpperCAmelCase__							:			Optional[int]	):
    __lowercase								=						TFDistilBertModel(config=UpperCAmelCase__	)
    __lowercase								=						{"input_ids": input_ids, "attention_mask": input_mask}

    __lowercase								=						model(UpperCAmelCase__	)

    __lowercase								=						[input_ids, input_mask]

    __lowercase								=						model(UpperCAmelCase__	)

    self.parent.assertEqual(result.last_hidden_state.shape,					(self.batch_size, self.seq_length, self.hidden_size)	)


 def 					_lowercase (			self							:			Tuple,					UpperCAmelCase__							:			Optional[Any],					UpperCAmelCase__							:			Any,					UpperCAmelCase__							:			Tuple,					UpperCAmelCase__							:			Any,					UpperCAmelCase__							:			Union[str, Any],					UpperCAmelCase__							:			List[str]	):
    __lowercase								=						TFDistilBertForMaskedLM(config=UpperCAmelCase__	)
    __lowercase								=						{"input_ids": input_ids, "attention_mask": input_mask}
    __lowercase								=						model(UpperCAmelCase__	)
    self.parent.assertEqual(result.logits.shape,					(self.batch_size, self.seq_length, self.vocab_size)	)


 def 					_lowercase (			self							:			int,					UpperCAmelCase__							:			Optional[Any],					UpperCAmelCase__							:			List[str],					UpperCAmelCase__							:			Dict,					UpperCAmelCase__							:			Any,					UpperCAmelCase__							:			Dict,					UpperCAmelCase__							:			Any	):
    __lowercase								=						TFDistilBertForQuestionAnswering(config=UpperCAmelCase__	)
    __lowercase								=						{
        "input_ids": input_ids,
        "attention_mask": input_mask,
    }
    __lowercase								=						model(UpperCAmelCase__	)
    self.parent.assertEqual(result.start_logits.shape,					(self.batch_size, self.seq_length)	)
    self.parent.assertEqual(result.end_logits.shape,					(self.batch_size, self.seq_length)	)


 def 					_lowercase (			self							:			Any,					UpperCAmelCase__							:			Optional[int],					UpperCAmelCase__							:			Any,					UpperCAmelCase__							:			Optional[int],					UpperCAmelCase__							:			List[Any],					UpperCAmelCase__							:			List[str],					UpperCAmelCase__							:			Any	):
    __lowercase								=						self.num_labels
    __lowercase								=						TFDistilBertForSequenceClassification(UpperCAmelCase__	)
    __lowercase								=						{"input_ids": input_ids, "attention_mask": input_mask}
    __lowercase								=						model(UpperCAmelCase__	)
    self.parent.assertEqual(result.logits.shape,					(self.batch_size, self.num_labels)	)


 def 					_lowercase (			self							:			Optional[int],					UpperCAmelCase__							:			str,					UpperCAmelCase__							:			Optional[int],					UpperCAmelCase__							:			Dict,					UpperCAmelCase__							:			int,					UpperCAmelCase__							:			List[Any],					UpperCAmelCase__							:			Tuple	):
    __lowercase								=						self.num_choices
    __lowercase								=						TFDistilBertForMultipleChoice(UpperCAmelCase__	)
    __lowercase								=						tf.tile(tf.expand_dims(UpperCAmelCase__,					1	),					(1, self.num_choices, 1)	)
    __lowercase								=						tf.tile(tf.expand_dims(UpperCAmelCase__,					1	),					(1, self.num_choices, 1)	)
    __lowercase								=						{
        "input_ids": multiple_choice_inputs_ids,
        "attention_mask": multiple_choice_input_mask,
    }
    __lowercase								=						model(UpperCAmelCase__	)
    self.parent.assertEqual(result.logits.shape,					(self.batch_size, self.num_choices)	)


 def 					_lowercase (			self							:			str,					UpperCAmelCase__							:			Union[str, Any],					UpperCAmelCase__							:			str,					UpperCAmelCase__							:			Any,					UpperCAmelCase__							:			Optional[int],					UpperCAmelCase__							:			int,					UpperCAmelCase__							:			List[Any]	):
    __lowercase								=						self.num_labels
    __lowercase								=						TFDistilBertForTokenClassification(UpperCAmelCase__	)
    __lowercase								=						{"input_ids": input_ids, "attention_mask": input_mask}
    __lowercase								=						model(UpperCAmelCase__	)
    self.parent.assertEqual(result.logits.shape,					(self.batch_size, self.seq_length, self.num_labels)	)



 def 					_lowercase (			self							:			Dict	):
    __lowercase								=						self.prepare_config_and_inputs()
    ((__lowercase)		,(__lowercase)		,(__lowercase)		,(__lowercase)		,(__lowercase)		,(__lowercase))								=						config_and_inputs
    __lowercase								=						{"input_ids": input_ids, "attention_mask": input_mask}
    return config, inputs_dict







@require_tf
class 					_lowerCAmelCase  (	lowercase     ,lowercase     ,unittest.TestCase   ):







 """simple docstring"""







 __UpperCAmelCase	:					int   		=	(
     (
         TFDistilBertModel,
         TFDistilBertForMaskedLM,
         TFDistilBertForQuestionAnswering,
         TFDistilBertForSequenceClassification,
         TFDistilBertForTokenClassification,
         TFDistilBertForMultipleChoice,
     )
     if is_tf_available()
     else None
 )
 __UpperCAmelCase	:					Union[str, Any]   		=	(
     {
         "feature-extraction": TFDistilBertModel,
         "fill-mask": TFDistilBertForMaskedLM,
         "question-answering": TFDistilBertForQuestionAnswering,
         "text-classification": TFDistilBertForSequenceClassification,
         "token-classification": TFDistilBertForTokenClassification,
         "zero-shot": TFDistilBertForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 __UpperCAmelCase	:					str   		=	False
 __UpperCAmelCase	:					List[str]   		=	False


 def 					_lowercase (			self							:			Optional[int]	):
    __lowercase								=						TFDistilBertModelTester(self	)
    __lowercase								=						ConfigTester(self,					config_class=UpperCAmelCase__,					dim=3_7	)


 def 					_lowercase (			self							:			str	):
    self.config_tester.run_common_tests()


 def 					_lowercase (			self							:			Optional[int]	):
    __lowercase								=						self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_distilbert_model(*UpperCAmelCase__	)


 def 					_lowercase (			self							:			Dict	):
    __lowercase								=						self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_distilbert_for_masked_lm(*UpperCAmelCase__	)


 def 					_lowercase (			self							:			List[Any]	):
    __lowercase								=						self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_distilbert_for_question_answering(*UpperCAmelCase__	)


 def 					_lowercase (			self							:			Dict	):
    __lowercase								=						self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_distilbert_for_sequence_classification(*UpperCAmelCase__	)


 def 					_lowercase (			self							:			List[str]	):
    __lowercase								=						self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_distilbert_for_multiple_choice(*UpperCAmelCase__	)


 def 					_lowercase (			self							:			List[str]	):
    __lowercase								=						self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_distilbert_for_token_classification(*UpperCAmelCase__	)



 @slow
 def 					_lowercase (			self							:			str	):
    for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]	):
       __lowercase								=						TFDistilBertModel.from_pretrained(UpperCAmelCase__	)
       self.assertIsNotNone(UpperCAmelCase__	)







@require_tf
class 					_lowerCAmelCase  (	unittest.TestCase   ):







 """simple docstring"""







 @slow
 def 					_lowercase (			self							:			Union[str, Any]	):
    __lowercase								=						TFDistilBertModel.from_pretrained("distilbert-base-uncased"	)
    __lowercase								=						tf.constant([[0, 1, 2, 3, 4, 5]]	)
    __lowercase								=						model(UpperCAmelCase__	)[0]

    __lowercase								=						[1, 6, 7_6_8]
    self.assertEqual(output.shape,					UpperCAmelCase__	)

    __lowercase								=						tf.constant(
        [
            [
                [0.19_261_885, -0.13_732_955, 0.4_119_799],
                [0.22_150_156, -0.07_422_661, 0.39_037_204],
                [0.22_756_018, -0.0_896_414, 0.3_701_467],
            ]
        ]	)
    tf.debugging.assert_near(output[:, :3, :3],					UpperCAmelCase__,					atol=1E-4	)








"""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       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       def __init__(						self		,     snake_case__=0.01		,     snake_case__=1_000						):






         """simple docstring"""





         lowerCAmelCase     :     List[Any]								=							p_stop
         lowerCAmelCase     :     Optional[Any]								=							max_length







       def __iter__(						self						):






         """simple docstring"""





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


class       SCREAMING_SNAKE_CASE__				(		unittest.TestCase				):
       """simple docstring"""


       def        lowercase__							(						self		,     snake_case__		,     snake_case__		,     snake_case__=False		,     snake_case__=True						):






         """simple docstring"""





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

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     Union[str, Any]								=							BatchSampler(range(24						)		,     batch_size=3		,     drop_last=snake_case__						)
         lowerCAmelCase     :     Any								=							[
             [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
             [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
         ]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__						)

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

         # Check the shards when the dataset is a round multiple of batch size but not total batch size.
         lowerCAmelCase     :     Union[str, Any]								=							BatchSampler(range(21						)		,     batch_size=3		,     drop_last=snake_case__						)
         lowerCAmelCase     :     List[str]								=							[
             [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
             [[3, 4, 5], [9, 10, 11], [15, 16, 17], [0, 1, 2]],
         ]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__						)

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

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

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

         # Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
         # num_processes batch.
         lowerCAmelCase     :     Any								=							BatchSampler(range(20						)		,     batch_size=3		,     drop_last=snake_case__						)
         lowerCAmelCase     :     Tuple								=							[
             [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 0]],
             [[3, 4, 5], [9, 10, 11], [15, 16, 17], [1, 2, 3]],
         ]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__						)

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

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

         lowerCAmelCase     :     Optional[Any]								=							BatchSampler(range(2						)		,     batch_size=3		,     drop_last=snake_case__						)
         lowerCAmelCase     :     Any								=							[[], []]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__						)

       def        lowercase__							(						self						):






         """simple docstring"""





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

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

         # Check the shards when the dataset is not a round multiple of batch size.
         lowerCAmelCase     :     Optional[int]								=							BatchSampler(range(22						)		,     batch_size=4		,     drop_last=snake_case__						)
         lowerCAmelCase     :     Optional[Any]								=							[
             [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
             [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [0, 1]],
         ]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     split_batches=snake_case__						)

         lowerCAmelCase     :     List[Any]								=							BatchSampler(range(22						)		,     batch_size=4		,     drop_last=snake_case__						)
         lowerCAmelCase     :     List[str]								=							[
             [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
             [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
         ]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     split_batches=snake_case__						)

         # Check the shards when the dataset is not a round multiple of batch size or num_processes.
         lowerCAmelCase     :     Tuple								=							BatchSampler(range(21						)		,     batch_size=4		,     drop_last=snake_case__						)
         lowerCAmelCase     :     List[str]								=							[
             [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 0]],
             [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [1, 2]],
         ]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     split_batches=snake_case__						)

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

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

         lowerCAmelCase     :     Optional[Any]								=							BatchSampler(range(2						)		,     batch_size=4		,     drop_last=snake_case__						)
         lowerCAmelCase     :     Optional[int]								=							[[], []]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     split_batches=snake_case__						)

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     Optional[int]								=							BatchSampler(range(24						)		,     batch_size=3		,     drop_last=snake_case__						)
         lowerCAmelCase     :     Any								=							[
             [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
             [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
         ]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     even_batches=snake_case__						)

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

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

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

         # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
         # num_processes batch.
         lowerCAmelCase     :     str								=							BatchSampler(range(22						)		,     batch_size=3		,     drop_last=snake_case__						)
         lowerCAmelCase     :     Union[str, Any]								=							[
             [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
             [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21]],
         ]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     even_batches=snake_case__						)

         lowerCAmelCase     :     List[str]								=							BatchSampler(range(22						)		,     batch_size=3		,     drop_last=snake_case__						)
         lowerCAmelCase     :     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(snake_case__		,     snake_case__		,     even_batches=snake_case__						)

         # Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
         # num_processes batch.
         lowerCAmelCase     :     str								=							BatchSampler(range(20						)		,     batch_size=3		,     drop_last=snake_case__						)
         lowerCAmelCase     :     Union[str, Any]								=							[
             [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19]],
             [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
         ]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     even_batches=snake_case__						)

         lowerCAmelCase     :     Any								=							BatchSampler(range(20						)		,     batch_size=3		,     drop_last=snake_case__						)
         lowerCAmelCase     :     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(snake_case__		,     snake_case__		,     even_batches=snake_case__						)

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

         lowerCAmelCase     :     str								=							BatchSampler(range(2						)		,     batch_size=3		,     drop_last=snake_case__						)
         lowerCAmelCase     :     Optional[Any]								=							[[], []]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     even_batches=snake_case__						)

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     str								=							BatchSampler(range(24						)		,     batch_size=4		,     drop_last=snake_case__						)
         lowerCAmelCase     :     Dict								=							[
             [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
             [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
         ]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     split_batches=snake_case__		,     even_batches=snake_case__						)

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

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

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

         # Check the shards when the dataset is not a round multiple of batch size or num_processes.
         lowerCAmelCase     :     int								=							BatchSampler(range(21						)		,     batch_size=4		,     drop_last=snake_case__						)
         lowerCAmelCase     :     str								=							[
             [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20]],
             [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
         ]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     split_batches=snake_case__		,     even_batches=snake_case__						)

         lowerCAmelCase     :     str								=							BatchSampler(range(21						)		,     batch_size=4		,     drop_last=snake_case__						)
         lowerCAmelCase     :     Union[str, Any]								=							[
             [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
             [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
         ]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     split_batches=snake_case__		,     even_batches=snake_case__						)

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

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

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     List[str]								=							[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
         lowerCAmelCase     :     Tuple								=							[BatchSamplerShard(snake_case__		,     2		,     snake_case__		,     even_batches=snake_case__						) for i in range(2						)]

         self.assertEqual(len(batch_sampler_shards[0]						)		,     3						)
         self.assertEqual(len(batch_sampler_shards[1]						)		,     2						)

         self.assertListEqual(list(batch_sampler_shards[0]						)		,     [[0, 1, 2], [5, 6, 7, 8], [12, 13]]						)
         self.assertListEqual(list(batch_sampler_shards[1]						)		,     [[3, 4], [9, 10, 11]]						)

       def        lowercase__							(						self		,     snake_case__		,     snake_case__		,     snake_case__		,     snake_case__=False		,     snake_case__=2		,     snake_case__=False						):






         """simple docstring"""





         random.seed(snake_case__						)
         lowerCAmelCase     :     List[str]								=							list(snake_case__						)

         lowerCAmelCase     :     Optional[int]								=							[
             IterableDatasetShard(
                 snake_case__		,     batch_size=snake_case__		,     drop_last=snake_case__		,     num_processes=snake_case__		,     process_index=snake_case__		,     split_batches=snake_case__		,     )
             for i in range(snake_case__						)
         ]
         lowerCAmelCase     :     str								=							[]
         for iterable_dataset_shard in iterable_dataset_shards:
           # Since our random iterable dataset will be... random... we need to use a seed to get reproducible results.
           random.seed(snake_case__						)
           iterable_dataset_lists.append(list(snake_case__						)						)

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

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

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

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     List[str]								=							42
         lowerCAmelCase     :     Tuple								=							RandomIterableDataset()

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

         # Edge case with a very small dataset
         lowerCAmelCase     :     List[str]								=							RandomIterableDataset(max_length=2						)

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

       def        lowercase__							(						self						):






         """simple docstring"""





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

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     List[str]								=							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						):






         """simple docstring"""





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

       def        lowercase__							(						self						):






         """simple docstring"""





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

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







       def        lowercase__							(						self						):






         """simple docstring"""





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

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





