



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







							def __init__(     self,    __a=0.01,    __a=1000):

														'''simple docstring'''

														_lowerCAmelCase  :							Optional[Any]				     =					p_stop
														_lowerCAmelCase  :							Dict				     =					max_length







							def __iter__(     self):

														'''simple docstring'''

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



class 						UpperCAmelCase_					( unittest.TestCase):







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

														'''simple docstring'''

														_lowerCAmelCase  :							Dict				     =					[
														    BatchSamplerShard(__a,    2,    __a,    split_batches=__a,    even_batches=__a)
														    for i in range(2)
														]
														_lowerCAmelCase  :							Optional[Any]				     =					[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      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase  :							Optional[Any]				     =					BatchSampler(range(24),    batch_size=3,    drop_last=__a)
														_lowerCAmelCase  :							Dict				     =					[
														    [[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)

														_lowerCAmelCase  :							str				     =					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.
														_lowerCAmelCase  :							Union[str, Any]				     =					BatchSampler(range(21),    batch_size=3,    drop_last=__a)
														_lowerCAmelCase  :							Optional[Any]				     =					[
														    [[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)

														_lowerCAmelCase  :							Union[str, Any]				     =					BatchSampler(range(21),    batch_size=3,    drop_last=__a)
														_lowerCAmelCase  :							int				     =					[
														    [[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.
														_lowerCAmelCase  :							List[Any]				     =					BatchSampler(range(22),    batch_size=3,    drop_last=__a)
														_lowerCAmelCase  :							int				     =					[
														    [[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)

														_lowerCAmelCase  :							Tuple				     =					BatchSampler(range(22),    batch_size=3,    drop_last=__a)
														_lowerCAmelCase  :							Optional[int]				     =					[
														    [[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.
														_lowerCAmelCase  :							int				     =					BatchSampler(range(20),    batch_size=3,    drop_last=__a)
														_lowerCAmelCase  :							List[Any]				     =					[
														    [[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)

														_lowerCAmelCase  :							List[Any]				     =					BatchSampler(range(20),    batch_size=3,    drop_last=__a)
														_lowerCAmelCase  :							int				     =					[
														    [[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.
														_lowerCAmelCase  :							str				     =					BatchSampler(range(2),    batch_size=3,    drop_last=__a)
														_lowerCAmelCase  :							str				     =					[[[0, 1, 0]], [[1, 0, 1]]]
														self.check_batch_sampler_shards(__a,    __a)

														_lowerCAmelCase  :							Union[str, Any]				     =					BatchSampler(range(2),    batch_size=3,    drop_last=__a)
														_lowerCAmelCase  :							Dict				     =					[[], []]
														self.check_batch_sampler_shards(__a,    __a)







							def      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase  :							Dict				     =					BatchSampler(range(24),    batch_size=4,    drop_last=__a)
														_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], [22, 23]],
														]
														self.check_batch_sampler_shards(__a,    __a,    split_batches=__a)

														_lowerCAmelCase  :							Optional[int]				     =					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.
														_lowerCAmelCase  :							Union[str, Any]				     =					BatchSampler(range(22),    batch_size=4,    drop_last=__a)
														_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], [0, 1]],
														]
														self.check_batch_sampler_shards(__a,    __a,    split_batches=__a)

														_lowerCAmelCase  :							List[str]				     =					BatchSampler(range(22),    batch_size=4,    drop_last=__a)
														_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(__a,    __a,    split_batches=__a)

														# Check the shards when the dataset is not a round multiple of batch size or num_processes.
														_lowerCAmelCase  :							str				     =					BatchSampler(range(21),    batch_size=4,    drop_last=__a)
														_lowerCAmelCase  :							List[Any]				     =					[
														    [[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)

														_lowerCAmelCase  :							int				     =					BatchSampler(range(21),    batch_size=4,    drop_last=__a)
														_lowerCAmelCase  :							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(__a,    __a,    split_batches=__a)

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

														_lowerCAmelCase  :							Any				     =					BatchSampler(range(2),    batch_size=4,    drop_last=__a)
														_lowerCAmelCase  :							Dict				     =					[[], []]
														self.check_batch_sampler_shards(__a,    __a,    split_batches=__a)







							def      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase  :							Tuple				     =					BatchSampler(range(24),    batch_size=3,    drop_last=__a)
														_lowerCAmelCase  :							Optional[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(__a,    __a,    even_batches=__a)

														_lowerCAmelCase  :							Any				     =					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.
														_lowerCAmelCase  :							List[Any]				     =					BatchSampler(range(21),    batch_size=3,    drop_last=__a)
														_lowerCAmelCase  :							Optional[int]				     =					[
														    [[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)

														_lowerCAmelCase  :							List[str]				     =					BatchSampler(range(21),    batch_size=3,    drop_last=__a)
														_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(__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.
														_lowerCAmelCase  :							Tuple				     =					BatchSampler(range(22),    batch_size=3,    drop_last=__a)
														_lowerCAmelCase  :							int				     =					[
														    [[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)

														_lowerCAmelCase  :							List[str]				     =					BatchSampler(range(22),    batch_size=3,    drop_last=__a)
														_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(__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.
														_lowerCAmelCase  :							Union[str, Any]				     =					BatchSampler(range(20),    batch_size=3,    drop_last=__a)
														_lowerCAmelCase  :							Dict				     =					[
														    [[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)

														_lowerCAmelCase  :							Dict				     =					BatchSampler(range(20),    batch_size=3,    drop_last=__a)
														_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(__a,    __a,    even_batches=__a)

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

														_lowerCAmelCase  :							Tuple				     =					BatchSampler(range(2),    batch_size=3,    drop_last=__a)
														_lowerCAmelCase  :							Tuple				     =					[[], []]
														self.check_batch_sampler_shards(__a,    __a,    even_batches=__a)







							def      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase  :							Tuple				     =					BatchSampler(range(24),    batch_size=4,    drop_last=__a)
														_lowerCAmelCase  :							Tuple				     =					[
														    [[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)

														_lowerCAmelCase  :							Union[str, Any]				     =					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.
														_lowerCAmelCase  :							List[Any]				     =					BatchSampler(range(22),    batch_size=4,    drop_last=__a)
														_lowerCAmelCase  :							Optional[int]				     =					[
														    [[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)

														_lowerCAmelCase  :							List[Any]				     =					BatchSampler(range(22),    batch_size=4,    drop_last=__a)
														_lowerCAmelCase  :							Optional[int]				     =					[
														    [[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.
														_lowerCAmelCase  :							Any				     =					BatchSampler(range(21),    batch_size=4,    drop_last=__a)
														_lowerCAmelCase  :							Optional[Any]				     =					[
														    [[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)

														_lowerCAmelCase  :							Tuple				     =					BatchSampler(range(21),    batch_size=4,    drop_last=__a)
														_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(__a,    __a,    split_batches=__a,    even_batches=__a)

														# Check the shards when the dataset is very small.
														_lowerCAmelCase  :							Dict				     =					BatchSampler(range(2),    batch_size=4,    drop_last=__a)
														_lowerCAmelCase  :							List[str]				     =					[[[0, 1]], []]
														self.check_batch_sampler_shards(__a,    __a,    split_batches=__a,    even_batches=__a)

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







							def      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase  :							int				     =					[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
														_lowerCAmelCase  :							List[Any]				     =					[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      snake_case__   (     self,    __a,    __a,    __a,    __a=False,    __a=2,    __a=False):

														'''simple docstring'''

														random.seed(__a)
														_lowerCAmelCase  :							Dict				     =					list(__a)

														_lowerCAmelCase  :							str				     =					[
														    IterableDatasetShard(
														        __a,    batch_size=__a,    drop_last=__a,    num_processes=__a,    process_index=__a,    split_batches=__a,    )
														    for i in range(__a)
														]
														_lowerCAmelCase  :							Dict				     =					[]
														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))

														_lowerCAmelCase  :							str				     =					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  :							Any				     =					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)

														_lowerCAmelCase  :							Optional[int]				     =					[]
														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      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase  :							Tuple				     =					42
														_lowerCAmelCase  :							Dict				     =					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
														_lowerCAmelCase  :							List[str]				     =					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      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase  :							Dict				     =					BatchSampler(range(16),    batch_size=4,    drop_last=__a)
														_lowerCAmelCase  :							Dict				     =					SkipBatchSampler(__a,    2)
														self.assertListEqual(list(__a),    [[8, 9, 10, 11], [12, 13, 14, 15]])







							def      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase  :							Tuple				     =					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      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase  :							Tuple				     =					DataLoader(list(range(16)),    batch_size=4)
														_lowerCAmelCase  :							Dict				     =					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      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase  :							int				     =					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      snake_case__   (     self):

														'''simple docstring'''

														Accelerator()
														_lowerCAmelCase  :							Optional[int]				     =					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)





import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import CONFIG_MAPPING


_snake_case	=						logging.get_logger(__name__)



class 						UpperCAmelCase_					( a):
							lowerCamelCase__         =	'upernet'







							def __init__(     self,    __a=None,    __a=512,    __a=0.02,    __a=[1, 2, 3, 6],    __a=True,    __a=0.4,    __a=384,    __a=256,    __a=1,    __a=False,    __a=255,    **__a,    ):

														'''simple docstring'''

														super().__init__(**__a)

														if backbone_config is None:
																					logger.info("`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone.")
																					_lowerCAmelCase  :							List[str]				     =					CONFIG_MAPPING["resnet"](out_features=["stage1", "stage2", "stage3", "stage4"])
														elif isinstance(__a,    __a):
																					_lowerCAmelCase  :							List[Any]				     =					backbone_config.get("model_type")
																					_lowerCAmelCase  :							Dict				     =					CONFIG_MAPPING[backbone_model_type]
																					_lowerCAmelCase  :							Optional[Any]				     =					config_class.from_dict(__a)

														_lowerCAmelCase  :							Tuple				     =					backbone_config
														_lowerCAmelCase  :							List[Any]				     =					hidden_size
														_lowerCAmelCase  :							Union[str, Any]				     =					initializer_range
														_lowerCAmelCase  :							str				     =					pool_scales
														_lowerCAmelCase  :							List[str]				     =					use_auxiliary_head
														_lowerCAmelCase  :							Dict				     =					auxiliary_loss_weight
														_lowerCAmelCase  :							Tuple				     =					auxiliary_in_channels
														_lowerCAmelCase  :							Optional[Any]				     =					auxiliary_channels
														_lowerCAmelCase  :							str				     =					auxiliary_num_convs
														_lowerCAmelCase  :							Union[str, Any]				     =					auxiliary_concat_input
														_lowerCAmelCase  :							Dict				     =					loss_ignore_index







							def      snake_case__   (     self):

														'''simple docstring'''

														_lowerCAmelCase  :							Any				     =					copy.deepcopy(self.__dict__)
														_lowerCAmelCase  :							List[Any]				     =					self.backbone_config.to_dict()
														_lowerCAmelCase  :							Optional[Any]				     =					self.__class__.model_type
														return output
