

import copy
import inspect
import unittest

from transformers import AutoBackbone
from transformers.configuration_utils import PretrainedConfig
from transformers.testing_utils import require_timm, require_torch, torch_device
from transformers.utils.import_utils import is_torch_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor


if is_torch_available():
       import torch

       from transformers import TimmBackbone, TimmBackboneConfig

from ...test_pipeline_mixin import PipelineTesterMixin






class 	UpperCAmelCase__		:






  """simple docstring"""


  def __init__(							self						:		Optional[int]     ,       __lowerCamelCase						:		int     ,       __lowerCamelCase						:		Union[str, Any]=None     ,       __lowerCamelCase						:		Optional[int]=None     ,       __lowerCamelCase						:		Tuple=None     ,       __lowerCamelCase						:		str="resnet50"     ,       __lowerCamelCase						:		Optional[Any]=3     ,       __lowerCamelCase						:		Tuple=32     ,       __lowerCamelCase						:		List[Any]=3     ,       __lowerCamelCase						:		Dict=True     ,       __lowerCamelCase						:		Dict=True     ,       )		->							Tuple:
       SCREAMING_SNAKE_CASE__   =   parent
       SCREAMING_SNAKE_CASE__   =   out_indices if out_indices is not None else [4]
       SCREAMING_SNAKE_CASE__   =   stage_names
       SCREAMING_SNAKE_CASE__   =   out_features
       SCREAMING_SNAKE_CASE__   =   backbone
       SCREAMING_SNAKE_CASE__   =   batch_size
       SCREAMING_SNAKE_CASE__   =   image_size
       SCREAMING_SNAKE_CASE__   =   num_channels
       SCREAMING_SNAKE_CASE__   =   use_pretrained_backbone
       SCREAMING_SNAKE_CASE__   =   is_training






  def 				lowercase_		(							self						:		Dict )		->							Union[str, Any]:
       SCREAMING_SNAKE_CASE__   =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )
       SCREAMING_SNAKE_CASE__   =   self.get_config()

       return config, pixel_values






  def 				lowercase_		(							self						:		str )		->							str:
       return TimmBackboneConfig(
           image_size=self.image_size     ,       num_channels=self.num_channels     ,       out_features=self.out_features     ,       out_indices=self.out_indices     ,       stage_names=self.stage_names     ,       use_pretrained_backbone=self.use_pretrained_backbone     ,       backbone=self.backbone     ,       )






  def 				lowercase_		(							self						:		Optional[Any]     ,       __lowerCamelCase						:		Union[str, Any]     ,       __lowerCamelCase						:		Dict )		->							Optional[Any]:
       SCREAMING_SNAKE_CASE__   =   TimmBackbone(config=__lowerCamelCase )
       model.to(__lowerCamelCase )
       model.eval()
       with torch.no_grad():
            SCREAMING_SNAKE_CASE__   =   model(__lowerCamelCase )
       self.parent.assertEqual(
           result.feature_map[-1].shape     ,       (self.batch_size, model.channels[-1], 14, 14)     ,       )







  def 				lowercase_		(							self						:		Dict )		->							int:
       SCREAMING_SNAKE_CASE__   =   self.prepare_config_and_inputs()
       SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   config_and_inputs
       SCREAMING_SNAKE_CASE__   =   {'''pixel_values''': pixel_values}
       return config, inputs_dict






@require_torch
@require_timm
class 	UpperCAmelCase__		(	A__						,					A__						,					A__						,					unittest.TestCase	):






  """simple docstring"""


  a									= (TimmBackbone,) if is_torch_available() else ()
  a									= {"feature-extraction": TimmBackbone} if is_torch_available() else {}
  a									= False
  a									= False
  a									= False
  a									= False






  def 				lowercase_		(							self						:		str )		->							Optional[int]:
       SCREAMING_SNAKE_CASE__   =   TimmBackboneModelTester(self )
       SCREAMING_SNAKE_CASE__   =   ConfigTester(self     ,       config_class=__lowerCamelCase     ,       has_text_modality=__lowerCamelCase )






  def 				lowercase_		(							self						:		Union[str, Any] )		->							List[Any]:
       self.config_tester.create_and_test_config_to_json_string()
       self.config_tester.create_and_test_config_to_json_file()
       self.config_tester.create_and_test_config_from_and_save_pretrained()
       self.config_tester.create_and_test_config_with_num_labels()
       self.config_tester.check_config_can_be_init_without_params()
       self.config_tester.check_config_arguments_init()






  def 				lowercase_		(							self						:		Any )		->							Union[str, Any]:
       SCREAMING_SNAKE_CASE__   =   '''resnet18'''
       SCREAMING_SNAKE_CASE__   =   '''microsoft/resnet-18'''

       SCREAMING_SNAKE_CASE__   =   AutoBackbone.from_pretrained(__lowerCamelCase     ,       use_timm_backbone=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   AutoBackbone.from_pretrained(__lowerCamelCase )

       self.assertEqual(len(timm_model.out_features )     ,       len(transformers_model.out_features ) )
       self.assertEqual(len(timm_model.stage_names )     ,       len(transformers_model.stage_names ) )
       self.assertEqual(timm_model.channels     ,       transformers_model.channels )
       # Out indices are set to the last layer by default. For timm models, we don't know
       # the number of layers in advance, so we set it to (-1,), whereas for transformers
       # models, we set it to [len(stage_names) - 1] (kept for backward compatibility).
       self.assertEqual(timm_model.out_indices     ,       (-1,) )
       self.assertEqual(transformers_model.out_indices     ,       [len(timm_model.stage_names ) - 1] )

       SCREAMING_SNAKE_CASE__   =   AutoBackbone.from_pretrained(__lowerCamelCase     ,       use_timm_backbone=__lowerCamelCase     ,       out_indices=[1, 2, 3] )
       SCREAMING_SNAKE_CASE__   =   AutoBackbone.from_pretrained(__lowerCamelCase     ,       out_indices=[1, 2, 3] )

       self.assertEqual(timm_model.out_indices     ,       transformers_model.out_indices )
       self.assertEqual(len(timm_model.out_features )     ,       len(transformers_model.out_features ) )
       self.assertEqual(timm_model.channels     ,       transformers_model.channels )






  @unittest.skip('''TimmBackbone doesn\'t support feed forward chunking''' )
  def 				lowercase_		(							self						:		str )		->							int:
       pass






  @unittest.skip('''TimmBackbone doesn\'t have num_hidden_layers attribute''' )
  def 				lowercase_		(							self						:		Union[str, Any] )		->							int:
       pass






  @unittest.skip('''TimmBackbone initialization is managed on the timm side''' )
  def 				lowercase_		(							self						:		Any )		->							Dict:
       pass






  @unittest.skip('''TimmBackbone models doesn\'t have inputs_embeds''' )
  def 				lowercase_		(							self						:		int )		->							Dict:
       pass






  @unittest.skip('''TimmBackbone models doesn\'t have inputs_embeds''' )
  def 				lowercase_		(							self						:		List[Any] )		->							Optional[Any]:
       pass






  @unittest.skip('''TimmBackbone model cannot be created without specifying a backbone checkpoint''' )
  def 				lowercase_		(							self						:		Optional[Any] )		->							int:
       pass






  @unittest.skip('''Only checkpoints on timm can be loaded into TimmBackbone''' )
  def 				lowercase_		(							self						:		Any )		->							List[str]:
       pass






  @unittest.skip('''model weights aren\'t tied in TimmBackbone.''' )
  def 				lowercase_		(							self						:		Dict )		->							str:
       pass






  @unittest.skip('''model weights aren\'t tied in TimmBackbone.''' )
  def 				lowercase_		(							self						:		Optional[Any] )		->							int:
       pass






  @unittest.skip('''Only checkpoints on timm can be loaded into TimmBackbone''' )
  def 				lowercase_		(							self						:		Dict )		->							List[Any]:
       pass






  @unittest.skip('''Only checkpoints on timm can be loaded into TimmBackbone''' )
  def 				lowercase_		(							self						:		List[Any] )		->							str:
       pass






  @unittest.skip('''TimmBackbone doesn\'t have hidden size info in its configuration.''' )
  def 				lowercase_		(							self						:		List[str] )		->							List[Any]:
       pass






  @unittest.skip('''TimmBackbone doesn\'t support output_attentions.''' )
  def 				lowercase_		(							self						:		List[str] )		->							Optional[Any]:
       pass






  @unittest.skip('''Safetensors is not supported by timm.''' )
  def 				lowercase_		(							self						:		Tuple )		->							Optional[int]:
       pass






  @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.''' )
  def 				lowercase_		(							self						:		Any )		->							List[Any]:
       pass






  def 				lowercase_		(							self						:		Dict )		->							List[str]:
       SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
            SCREAMING_SNAKE_CASE__   =   model_class(__lowerCamelCase )
            SCREAMING_SNAKE_CASE__   =   inspect.signature(model.forward )
            # signature.parameters is an OrderedDict => so arg_names order is deterministic
            SCREAMING_SNAKE_CASE__   =   [*signature.parameters.keys()]

            SCREAMING_SNAKE_CASE__   =   ['''pixel_values''']
            self.assertListEqual(arg_names[:1]     ,       __lowerCamelCase )






  def 				lowercase_		(							self						:		Dict )		->							Optional[int]:
       SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   self.model_tester.prepare_config_and_inputs_for_common()
       SCREAMING_SNAKE_CASE__   =   True
       SCREAMING_SNAKE_CASE__   =   self.has_attentions

       # no need to test all models as different heads yield the same functionality
       SCREAMING_SNAKE_CASE__   =   self.all_model_classes[0]
       SCREAMING_SNAKE_CASE__   =   model_class(__lowerCamelCase )
       model.to(__lowerCamelCase )

       SCREAMING_SNAKE_CASE__   =   self._prepare_for_class(__lowerCamelCase     ,       __lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   model(**__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   outputs[0][-1]

       # Encoder-/Decoder-only models
       SCREAMING_SNAKE_CASE__   =   outputs.hidden_states[0]
       hidden_states.retain_grad()

       if self.has_attentions:
            SCREAMING_SNAKE_CASE__   =   outputs.attentions[0]
            attentions.retain_grad()

       output.flatten()[0].backward(retain_graph=__lowerCamelCase )

       self.assertIsNotNone(hidden_states.grad )

       if self.has_attentions:
            self.assertIsNotNone(attentions.grad )







  def 				lowercase_		(							self						:		Union[str, Any] )		->							int:
       SCREAMING_SNAKE_CASE__,SCREAMING_SNAKE_CASE__   =   self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
            SCREAMING_SNAKE_CASE__   =   model_class(__lowerCamelCase )
            model.to(__lowerCamelCase )
            model.eval()
            SCREAMING_SNAKE_CASE__   =   model(**__lowerCamelCase )

            self.assertEqual(len(result.feature_maps )     ,       len(config.out_indices ) )
            self.assertEqual(len(model.channels )     ,       len(config.out_indices ) )

            # Check output of last stage is taken if out_features=None, out_indices=None
            SCREAMING_SNAKE_CASE__   =   copy.deepcopy(__lowerCamelCase )
            SCREAMING_SNAKE_CASE__   =   None
            SCREAMING_SNAKE_CASE__   =   model_class(__lowerCamelCase )
            model.to(__lowerCamelCase )
            model.eval()
            SCREAMING_SNAKE_CASE__   =   model(**__lowerCamelCase )

            self.assertEqual(len(result.feature_maps )     ,       1 )
            self.assertEqual(len(model.channels )     ,       1 )

            # Check backbone can be initialized with fresh weights
            SCREAMING_SNAKE_CASE__   =   copy.deepcopy(__lowerCamelCase )
            SCREAMING_SNAKE_CASE__   =   False
            SCREAMING_SNAKE_CASE__   =   model_class(__lowerCamelCase )
            model.to(__lowerCamelCase )
            model.eval()
            SCREAMING_SNAKE_CASE__   =   model(**__lowerCamelCase )





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						:		List[str]     ,       __lowerCamelCase						:		Tuple=0.01     ,       __lowerCamelCase						:		Optional[Any]=1000 )		->							Union[str, Any]:
       SCREAMING_SNAKE_CASE__   =   p_stop
       SCREAMING_SNAKE_CASE__   =   max_length







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






class 	UpperCAmelCase__		(	unittest.TestCase	):






  """simple docstring"""


  def 				lowercase_		(							self						:		List[Any]     ,       __lowerCamelCase						:		Union[str, Any]     ,       __lowerCamelCase						:		str     ,       __lowerCamelCase						:		Any=False     ,       __lowerCamelCase						:		Optional[int]=True )		->							Dict:
       SCREAMING_SNAKE_CASE__   =   [
           BatchSamplerShard(__lowerCamelCase     ,       2     ,       __lowerCamelCase     ,       split_batches=__lowerCamelCase     ,       even_batches=__lowerCamelCase )
           for i in range(2 )
       ]
       SCREAMING_SNAKE_CASE__   =   [list(__lowerCamelCase ) for batch_sampler_shard in batch_sampler_shards]
       if not split_batches:
            self.assertListEqual([len(__lowerCamelCase ) for shard in batch_sampler_shards]     ,       [len(__lowerCamelCase ) for e in expected] )
       self.assertListEqual(__lowerCamelCase     ,       __lowerCamelCase )






  def 				lowercase_		(							self						:		Any )		->							Optional[Any]:
       # Check the shards when the dataset is a round multiple of total batch size.
       SCREAMING_SNAKE_CASE__   =   BatchSampler(range(24 )     ,       batch_size=3     ,       drop_last=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   [
           [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
           [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
       ]
       self.check_batch_sampler_shards(__lowerCamelCase     ,       __lowerCamelCase )

       SCREAMING_SNAKE_CASE__   =   BatchSampler(range(24 )     ,       batch_size=3     ,       drop_last=__lowerCamelCase )
       # Expected shouldn't change
       self.check_batch_sampler_shards(__lowerCamelCase     ,       __lowerCamelCase )

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

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

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

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

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

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

       # Check the shards when the dataset is very small.
       SCREAMING_SNAKE_CASE__   =   BatchSampler(range(2 )     ,       batch_size=3     ,       drop_last=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   [[[0, 1, 0]], [[1, 0, 1]]]
       self.check_batch_sampler_shards(__lowerCamelCase     ,       __lowerCamelCase )

       SCREAMING_SNAKE_CASE__   =   BatchSampler(range(2 )     ,       batch_size=3     ,       drop_last=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   [[], []]
       self.check_batch_sampler_shards(__lowerCamelCase     ,       __lowerCamelCase )






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

       SCREAMING_SNAKE_CASE__   =   BatchSampler(range(24 )     ,       batch_size=4     ,       drop_last=__lowerCamelCase )
       # Expected shouldn't change
       self.check_batch_sampler_shards(__lowerCamelCase     ,       __lowerCamelCase     ,       split_batches=__lowerCamelCase )

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

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

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

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

       # Check the shards when the dataset is very small.
       SCREAMING_SNAKE_CASE__   =   BatchSampler(range(2 )     ,       batch_size=4     ,       drop_last=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   [[[0, 1]], [[0, 1]]]
       self.check_batch_sampler_shards(__lowerCamelCase     ,       __lowerCamelCase     ,       split_batches=__lowerCamelCase )

       SCREAMING_SNAKE_CASE__   =   BatchSampler(range(2 )     ,       batch_size=4     ,       drop_last=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   [[], []]
       self.check_batch_sampler_shards(__lowerCamelCase     ,       __lowerCamelCase     ,       split_batches=__lowerCamelCase )






  def 				lowercase_		(							self						:		str )		->							Dict:
       # Check the shards when the dataset is a round multiple of total batch size.
       SCREAMING_SNAKE_CASE__   =   BatchSampler(range(24 )     ,       batch_size=3     ,       drop_last=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   [
           [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
           [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
       ]
       self.check_batch_sampler_shards(__lowerCamelCase     ,       __lowerCamelCase     ,       even_batches=__lowerCamelCase )

       SCREAMING_SNAKE_CASE__   =   BatchSampler(range(24 )     ,       batch_size=3     ,       drop_last=__lowerCamelCase )
       # Expected shouldn't change
       self.check_batch_sampler_shards(__lowerCamelCase     ,       __lowerCamelCase     ,       even_batches=__lowerCamelCase )

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

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

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

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

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

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

       # Check the shards when the dataset is very small.
       SCREAMING_SNAKE_CASE__   =   BatchSampler(range(2 )     ,       batch_size=3     ,       drop_last=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   [[[0, 1]], []]
       self.check_batch_sampler_shards(__lowerCamelCase     ,       __lowerCamelCase     ,       even_batches=__lowerCamelCase )

       SCREAMING_SNAKE_CASE__   =   BatchSampler(range(2 )     ,       batch_size=3     ,       drop_last=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   [[], []]
       self.check_batch_sampler_shards(__lowerCamelCase     ,       __lowerCamelCase     ,       even_batches=__lowerCamelCase )






  def 				lowercase_		(							self						:		List[str] )		->							Tuple:
       # Check the shards when the dataset is a round multiple of batch size.
       SCREAMING_SNAKE_CASE__   =   BatchSampler(range(24 )     ,       batch_size=4     ,       drop_last=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   [
           [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
           [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
       ]
       self.check_batch_sampler_shards(__lowerCamelCase     ,       __lowerCamelCase     ,       split_batches=__lowerCamelCase     ,       even_batches=__lowerCamelCase )

       SCREAMING_SNAKE_CASE__   =   BatchSampler(range(24 )     ,       batch_size=4     ,       drop_last=__lowerCamelCase )
       # Expected shouldn't change
       self.check_batch_sampler_shards(__lowerCamelCase     ,       __lowerCamelCase     ,       split_batches=__lowerCamelCase     ,       even_batches=__lowerCamelCase )

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

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

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

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

       # Check the shards when the dataset is very small.
       SCREAMING_SNAKE_CASE__   =   BatchSampler(range(2 )     ,       batch_size=4     ,       drop_last=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   [[[0, 1]], []]
       self.check_batch_sampler_shards(__lowerCamelCase     ,       __lowerCamelCase     ,       split_batches=__lowerCamelCase     ,       even_batches=__lowerCamelCase )

       SCREAMING_SNAKE_CASE__   =   BatchSampler(range(2 )     ,       batch_size=4     ,       drop_last=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   [[], []]
       self.check_batch_sampler_shards(__lowerCamelCase     ,       __lowerCamelCase     ,       split_batches=__lowerCamelCase     ,       even_batches=__lowerCamelCase )






  def 				lowercase_		(							self						:		Optional[int] )		->							List[str]:
       SCREAMING_SNAKE_CASE__   =   [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
       SCREAMING_SNAKE_CASE__   =   [BatchSamplerShard(__lowerCamelCase     ,       2     ,       __lowerCamelCase     ,       even_batches=__lowerCamelCase ) for i in range(2 )]

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

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






  def 				lowercase_		(							self						:		Optional[Any]     ,       __lowerCamelCase						:		Any     ,       __lowerCamelCase						:		str     ,       __lowerCamelCase						:		str     ,       __lowerCamelCase						:		Tuple=False     ,       __lowerCamelCase						:		Optional[int]=2     ,       __lowerCamelCase						:		Dict=False )		->							str:
       random.seed(__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   list(__lowerCamelCase )

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

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

       SCREAMING_SNAKE_CASE__   =   []
       for idx in range(0     ,       len(__lowerCamelCase )     ,       __lowerCamelCase ):
            for l in iterable_dataset_lists:
                 observed += l[idx : idx + shard_batch_size]

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






  def 				lowercase_		(							self						:		str )		->							Dict:
       SCREAMING_SNAKE_CASE__   =   42
       SCREAMING_SNAKE_CASE__   =   RandomIterableDataset()

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

       # Edge case with a very small dataset
       SCREAMING_SNAKE_CASE__   =   RandomIterableDataset(max_length=2 )

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






  def 				lowercase_		(							self						:		Optional[Any] )		->							Any:
       SCREAMING_SNAKE_CASE__   =   BatchSampler(range(16 )     ,       batch_size=4     ,       drop_last=__lowerCamelCase )
       SCREAMING_SNAKE_CASE__   =   SkipBatchSampler(__lowerCamelCase     ,       2 )
       self.assertListEqual(list(__lowerCamelCase )     ,       [[8, 9, 10, 11], [12, 13, 14, 15]] )






  def 				lowercase_		(							self						:		Union[str, Any] )		->							str:
       SCREAMING_SNAKE_CASE__   =   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						:		Dict )		->							List[Any]:
       SCREAMING_SNAKE_CASE__   =   DataLoader(list(range(16 ) )     ,       batch_size=4 )
       SCREAMING_SNAKE_CASE__   =   skip_first_batches(__lowerCamelCase     ,       num_batches=2 )
       self.assertListEqual([t.tolist() for t in new_dataloader]     ,       [[8, 9, 10, 11], [12, 13, 14, 15]] )






  def 				lowercase_		(							self						:		List[Any] )		->							List[str]:
       SCREAMING_SNAKE_CASE__   =   DataLoaderShard(list(range(16 ) )     ,       batch_size=4 )
       for idx, _ in enumerate(__lowerCamelCase ):
            self.assertEqual(dataloader.end_of_dataloader     ,       idx == 3 )

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







  def 				lowercase_		(							self						:		Union[str, Any] )		->							str:
       Accelerator()
       SCREAMING_SNAKE_CASE__   =   DataLoaderDispatcher(range(16 )     ,       batch_size=4 )
       for idx, _ in enumerate(__lowerCamelCase ):
            self.assertEqual(dataloader.end_of_dataloader     ,       idx == 3 )

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


