





'''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	(						__UpperCamelCase			):



     '''simple docstring'''






     def __init__(					self						,   lowerCamelCase=0.01						,   lowerCamelCase=1_000	):
       _snake_case     		=				p_stop
       _snake_case     		=				max_length






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







class       __SCREAMING_SNAKE_CASE	(						unittest.TestCase			):



     '''simple docstring'''






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

     def 	UpperCamelCase(					self	):
       # Check the shards when the dataset is a round multiple of total batch size.
       _snake_case     		=				BatchSampler(range(24	)						,   batch_size=3						,   drop_last=__A	)
       _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(__A						,   __A	)

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

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

       _snake_case     		=				BatchSampler(range(21	)						,   batch_size=3						,   drop_last=__A	)
       _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(__A						,   __A	)

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

       _snake_case     		=				BatchSampler(range(22	)						,   batch_size=3						,   drop_last=__A	)
       _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(__A						,   __A	)

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

       _snake_case     		=				BatchSampler(range(20	)						,   batch_size=3						,   drop_last=__A	)
       _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(__A						,   __A	)

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

       _snake_case     		=				BatchSampler(range(2	)						,   batch_size=3						,   drop_last=__A	)
       _snake_case     		=				[[], []]
       self.check_batch_sampler_shards(__A						,   __A	)

     def 	UpperCamelCase(					self	):
       # Check the shards when the dataset is a round multiple of batch size.
       _snake_case     		=				BatchSampler(range(24	)						,   batch_size=4						,   drop_last=__A	)
       _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(__A						,   __A						,   split_batches=__A	)

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

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

       _snake_case     		=				BatchSampler(range(22	)						,   batch_size=4						,   drop_last=__A	)
       _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(__A						,   __A						,   split_batches=__A	)

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

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

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

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

     def 	UpperCamelCase(					self	):
       # Check the shards when the dataset is a round multiple of total batch size.
       _snake_case     		=				BatchSampler(range(24	)						,   batch_size=3						,   drop_last=__A	)
       _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(__A						,   __A						,   even_batches=__A	)

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

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

       _snake_case     		=				BatchSampler(range(21	)						,   batch_size=3						,   drop_last=__A	)
       _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(__A						,   __A						,   even_batches=__A	)

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

       _snake_case     		=				BatchSampler(range(22	)						,   batch_size=3						,   drop_last=__A	)
       _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(__A						,   __A						,   even_batches=__A	)

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

       _snake_case     		=				BatchSampler(range(20	)						,   batch_size=3						,   drop_last=__A	)
       _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(__A						,   __A						,   even_batches=__A	)

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

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

     def 	UpperCamelCase(					self	):
       # Check the shards when the dataset is a round multiple of batch size.
       _snake_case     		=				BatchSampler(range(24	)						,   batch_size=4						,   drop_last=__A	)
       _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(__A						,   __A						,   split_batches=__A						,   even_batches=__A	)

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

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

       _snake_case     		=				BatchSampler(range(22	)						,   batch_size=4						,   drop_last=__A	)
       _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(__A						,   __A						,   split_batches=__A						,   even_batches=__A	)

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

       _snake_case     		=				BatchSampler(range(21	)						,   batch_size=4						,   drop_last=__A	)
       _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(__A						,   __A						,   split_batches=__A						,   even_batches=__A	)

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

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

     def 	UpperCamelCase(					self	):
       _snake_case     		=				[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
       _snake_case     		=				[BatchSamplerShard(__A						,   2						,   __A						,   even_batches=__A	) for i in range(2	)]

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

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

     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase						,   lowerCamelCase=False						,   lowerCamelCase=2						,   lowerCamelCase=False	):
       random.seed(__A	)
       _snake_case     		=				list(__A	)

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

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

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

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

     def 	UpperCamelCase(					self	):
       _snake_case     		=				42
       _snake_case     		=				RandomIterableDataset()

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

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

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

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

     def 	UpperCamelCase(					self	):
       _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 	UpperCamelCase(					self	):
       _snake_case     		=				DataLoader(list(range(16	)	)						,   batch_size=4	)
       _snake_case     		=				skip_first_batches(__A						,   num_batches=2	)
       self.assertListEqual([t.tolist() for t in new_dataloader]						,   [[8, 9, 10, 11], [12, 13, 14, 15]]	)

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

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






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

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









'''simple docstring'''



from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


__magic_name__   :						Any    =			logging.get_logger(__name__)

__magic_name__   :						Dict    =			{
    """google/bit-50""": """https://huggingface.co/google/bit-50/resolve/main/config.json""",
}


class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase					,    __UpperCamelCase			):



     '''simple docstring'''






     UpperCAmelCase__       :						List[str]      =   '''bit'''
     UpperCAmelCase__       :						Union[str, Any]      =   ['''preactivation''', '''bottleneck''']
     UpperCAmelCase__       :						int      =   ['''SAME''', '''VALID''']

     def __init__(					self						,   lowerCamelCase=3						,   lowerCamelCase=64						,   lowerCamelCase=[256, 512, 1_024, 2_048]						,   lowerCamelCase=[3, 4, 6, 3]						,   lowerCamelCase="preactivation"						,   lowerCamelCase="relu"						,   lowerCamelCase=None						,   lowerCamelCase=32						,   lowerCamelCase=0.0						,   lowerCamelCase=False						,   lowerCamelCase=32						,   lowerCamelCase=1						,   lowerCamelCase=None						,   lowerCamelCase=None						,   **lowerCamelCase						,   ):
       super().__init__(**lowerCamelCase	)
       if layer_type not in self.layer_types:
         raise ValueError(F'''layer_type={layer_type} is not one of {",".join(self.layer_types	)}'''	)
       if global_padding is not None:
         if global_padding.upper() in self.supported_padding:
           _snake_case     		=				global_padding.upper()
         else:
           raise ValueError(F'''Padding strategy {global_padding} not supported'''	)
       _snake_case     		=				num_channels
       _snake_case     		=				embedding_size
       _snake_case     		=				hidden_sizes
       _snake_case     		=				depths
       _snake_case     		=				layer_type
       _snake_case     		=				hidden_act
       _snake_case     		=				global_padding
       _snake_case     		=				num_groups
       _snake_case     		=				drop_path_rate
       _snake_case     		=				embedding_dynamic_padding
       _snake_case     		=				output_stride
       _snake_case     		=				width_factor

       _snake_case     		=				["stem"] + [F'''stage{idx}''' for idx in range(1						,   len(lowerCamelCase	) + 1	)]
       _snake_case		,		_snake_case     		=				get_aligned_output_features_output_indices(
           out_features=lowerCamelCase						,   out_indices=lowerCamelCase						,   stage_names=self.stage_names	)


