
"""simple docstring"""





import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..bit import BitConfig


_a		       =							logging.get_logger(__name__)

_a		       =							{
    'Intel/dpt-large': 'https://huggingface.co/Intel/dpt-large/resolve/main/config.json',
    # See all DPT models at https://huggingface.co/models?filter=dpt
}


class A_	(lowercase__      ):




  '''simple docstring'''
  SCREAMING_SNAKE_CASE__			:						str       =		"""dpt"""







  def __init__(		self				,   lowercase_=768				,   lowercase_=12				,   lowercase_=12				,   lowercase_=3072				,   lowercase_="gelu"				,   lowercase_=0.0				,   lowercase_=0.0				,   lowercase_=0.02				,   lowercase_=1E-1_2				,   lowercase_=384				,   lowercase_=16				,   lowercase_=3				,   lowercase_=False				,   lowercase_=True				,   lowercase_=[2, 5, 8, 11]				,   lowercase_="project"				,   lowercase_=[4, 2, 1, 0.5]				,   lowercase_=[96, 192, 384, 768]				,   lowercase_=256				,   lowercase_=-1				,   lowercase_=False				,   lowercase_=True				,   lowercase_=0.4				,   lowercase_=255				,   lowercase_=0.1				,   lowercase_=[1, 1024, 24, 24]				,   lowercase_=[0, 1]				,   lowercase_=None				,   **lowercase_				,   ):

     """simple docstring"""
     super().__init__(**lowercase_  )

     UpperCAmelCase_					:					int			     =    hidden_size
     UpperCAmelCase_					:					Optional[int]			     =    is_hybrid

     if self.is_hybrid:
        if backbone_config is None:
           logger.info("Initializing the config with a `BiT` backbone."  )
           UpperCAmelCase_					:					List[Any]			     =    {
               "global_padding": "same",
               "layer_type": "bottleneck",
               "depths": [3, 4, 9],
               "out_features": ["stage1", "stage2", "stage3"],
               "embedding_dynamic_padding": True,
           }
           UpperCAmelCase_					:					Tuple			     =    BitConfig(**lowercase_  )
        elif isinstance(lowercase_				,   lowercase_  ):
           logger.info("Initializing the config with a `BiT` backbone."  )
           UpperCAmelCase_					:					int			     =    BitConfig(**lowercase_  )
        elif isinstance(lowercase_				,   lowercase_  ):
           UpperCAmelCase_					:					Union[str, Any]			     =    backbone_config
        else:
           raise ValueError(
               F"""backbone_config must be a dictionary or a `PretrainedConfig`, got {backbone_config.__class__}."""  )

        UpperCAmelCase_					:					List[Any]			     =    backbone_featmap_shape
        UpperCAmelCase_					:					str			     =    neck_ignore_stages

        if readout_type != "project":
           raise ValueError("Readout type must be 'project' when using `DPT-hybrid` mode."  )
     else:
        UpperCAmelCase_					:					Any			     =    None
        UpperCAmelCase_					:					Optional[Any]			     =    None
        UpperCAmelCase_					:					List[str]			     =    []

     UpperCAmelCase_					:					int			     =    num_hidden_layers
     UpperCAmelCase_					:					List[Any]			     =    num_attention_heads
     UpperCAmelCase_					:					Dict			     =    intermediate_size
     UpperCAmelCase_					:					List[Any]			     =    hidden_act
     UpperCAmelCase_					:					Dict			     =    hidden_dropout_prob
     UpperCAmelCase_					:					Any			     =    attention_probs_dropout_prob
     UpperCAmelCase_					:					Optional[Any]			     =    initializer_range
     UpperCAmelCase_					:					int			     =    layer_norm_eps
     UpperCAmelCase_					:					Optional[int]			     =    image_size
     UpperCAmelCase_					:					int			     =    patch_size
     UpperCAmelCase_					:					str			     =    num_channels
     UpperCAmelCase_					:					Union[str, Any]			     =    qkv_bias
     UpperCAmelCase_					:					Optional[Any]			     =    backbone_out_indices
     if readout_type not in ["ignore", "add", "project"]:
        raise ValueError("Readout_type must be one of ['ignore', 'add', 'project']"  )
     UpperCAmelCase_					:					Any			     =    readout_type
     UpperCAmelCase_					:					str			     =    reassemble_factors
     UpperCAmelCase_					:					str			     =    neck_hidden_sizes
     UpperCAmelCase_					:					Optional[int]			     =    fusion_hidden_size
     UpperCAmelCase_					:					str			     =    head_in_index
     UpperCAmelCase_					:					Optional[Any]			     =    use_batch_norm_in_fusion_residual
     # auxiliary head attributes (semantic segmentation)
     UpperCAmelCase_					:					Optional[Any]			     =    use_auxiliary_head
     UpperCAmelCase_					:					Any			     =    auxiliary_loss_weight
     UpperCAmelCase_					:					List[str]			     =    semantic_loss_ignore_index
     UpperCAmelCase_					:					Union[str, Any]			     =    semantic_classifier_dropout







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					int			     =    copy.deepcopy(self.__dict__  )

     if output["backbone_config"] is not None:
        UpperCAmelCase_					:					int			     =    self.backbone_config.to_dict()

     UpperCAmelCase_					:					str			     =    self.__class__.model_type
     return output






"""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						)





