





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       __A							(   _snake_case	):




      """simple docstring"""







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



             """simple docstring"""





             _lowerCamelCase				:  List[str]						=		p_stop
             _lowerCamelCase				:  Tuple						=		max_length



      def __iter__(		self):



             """simple docstring"""





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







class       __A							(   unittest.TestCase	):




      """simple docstring"""







      def 					__snake_case      (		self   ,				a__   ,				a__   ,				a__=False   ,				a__=True):



             """simple docstring"""





             _lowerCamelCase				:  str						=		[
                 BatchSamplerShard(UpperCamelCase__   ,				2   ,				UpperCamelCase__   ,				split_batches=UpperCamelCase__   ,				even_batches=UpperCamelCase__)
                 for i in range(2)
             ]
             _lowerCamelCase				:  List[str]						=		[list(UpperCamelCase__) for batch_sampler_shard in batch_sampler_shards]
             if not split_batches:
                    self.assertListEqual([len(UpperCamelCase__) for shard in batch_sampler_shards]   ,				[len(UpperCamelCase__) for e in expected])
             self.assertListEqual(UpperCamelCase__   ,				UpperCamelCase__)


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  int						=		BatchSampler(range(24)   ,				batch_size=3   ,				drop_last=UpperCamelCase__)
             _lowerCamelCase				:  List[str]						=		[
                 [[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(UpperCamelCase__   ,				UpperCamelCase__)

             _lowerCamelCase				:  List[Any]						=		BatchSampler(range(24)   ,				batch_size=3   ,				drop_last=UpperCamelCase__)
             # Expected shouldn't change
             self.check_batch_sampler_shards(UpperCamelCase__   ,				UpperCamelCase__)

             # Check the shards when the dataset is a round multiple of batch size but not total batch size.
             _lowerCamelCase				:  Optional[int]						=		BatchSampler(range(21)   ,				batch_size=3   ,				drop_last=UpperCamelCase__)
             _lowerCamelCase				:  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(UpperCamelCase__   ,				UpperCamelCase__)

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

             # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
             # num_processes batch.
             _lowerCamelCase				:  Any						=		BatchSampler(range(22)   ,				batch_size=3   ,				drop_last=UpperCamelCase__)
             _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, 0, 1]],
             ]
             self.check_batch_sampler_shards(UpperCamelCase__   ,				UpperCamelCase__)

             _lowerCamelCase				:  Union[str, Any]						=		BatchSampler(range(22)   ,				batch_size=3   ,				drop_last=UpperCamelCase__)
             _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(UpperCamelCase__   ,				UpperCamelCase__)

             # 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				:  Tuple						=		BatchSampler(range(20)   ,				batch_size=3   ,				drop_last=UpperCamelCase__)
             _lowerCamelCase				:  Optional[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(UpperCamelCase__   ,				UpperCamelCase__)

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

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

             _lowerCamelCase				:  Union[str, Any]						=		BatchSampler(range(2)   ,				batch_size=3   ,				drop_last=UpperCamelCase__)
             _lowerCamelCase				:  Optional[int]						=		[[], []]
             self.check_batch_sampler_shards(UpperCamelCase__   ,				UpperCamelCase__)


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  int						=		BatchSampler(range(24)   ,				batch_size=4   ,				drop_last=UpperCamelCase__)
             _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(UpperCamelCase__   ,				UpperCamelCase__   ,				split_batches=UpperCamelCase__)

             _lowerCamelCase				:  Any						=		BatchSampler(range(24)   ,				batch_size=4   ,				drop_last=UpperCamelCase__)
             # Expected shouldn't change
             self.check_batch_sampler_shards(UpperCamelCase__   ,				UpperCamelCase__   ,				split_batches=UpperCamelCase__)

             # Check the shards when the dataset is not a round multiple of batch size.
             _lowerCamelCase				:  Dict						=		BatchSampler(range(22)   ,				batch_size=4   ,				drop_last=UpperCamelCase__)
             _lowerCamelCase				:  Union[str, 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(UpperCamelCase__   ,				UpperCamelCase__   ,				split_batches=UpperCamelCase__)

             _lowerCamelCase				:  Optional[Any]						=		BatchSampler(range(22)   ,				batch_size=4   ,				drop_last=UpperCamelCase__)
             _lowerCamelCase				:  Dict						=		[
                 [[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(UpperCamelCase__   ,				UpperCamelCase__   ,				split_batches=UpperCamelCase__)

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

             _lowerCamelCase				:  List[Any]						=		BatchSampler(range(21)   ,				batch_size=4   ,				drop_last=UpperCamelCase__)
             _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(UpperCamelCase__   ,				UpperCamelCase__   ,				split_batches=UpperCamelCase__)

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

             _lowerCamelCase				:  str						=		BatchSampler(range(2)   ,				batch_size=4   ,				drop_last=UpperCamelCase__)
             _lowerCamelCase				:  int						=		[[], []]
             self.check_batch_sampler_shards(UpperCamelCase__   ,				UpperCamelCase__   ,				split_batches=UpperCamelCase__)


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  Tuple						=		BatchSampler(range(24)   ,				batch_size=3   ,				drop_last=UpperCamelCase__)
             _lowerCamelCase				:  Optional[int]						=		[
                 [[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(UpperCamelCase__   ,				UpperCamelCase__   ,				even_batches=UpperCamelCase__)

             _lowerCamelCase				:  int						=		BatchSampler(range(24)   ,				batch_size=3   ,				drop_last=UpperCamelCase__)
             # Expected shouldn't change
             self.check_batch_sampler_shards(UpperCamelCase__   ,				UpperCamelCase__   ,				even_batches=UpperCamelCase__)

             # Check the shards when the dataset is a round multiple of batch size but not total batch size.
             _lowerCamelCase				:  Optional[int]						=		BatchSampler(range(21)   ,				batch_size=3   ,				drop_last=UpperCamelCase__)
             _lowerCamelCase				:  str						=		[
                 [[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(UpperCamelCase__   ,				UpperCamelCase__   ,				even_batches=UpperCamelCase__)

             _lowerCamelCase				:  List[Any]						=		BatchSampler(range(21)   ,				batch_size=3   ,				drop_last=UpperCamelCase__)
             _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(UpperCamelCase__   ,				UpperCamelCase__   ,				even_batches=UpperCamelCase__)

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

             _lowerCamelCase				:  int						=		BatchSampler(range(22)   ,				batch_size=3   ,				drop_last=UpperCamelCase__)
             _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(UpperCamelCase__   ,				UpperCamelCase__   ,				even_batches=UpperCamelCase__)

             # 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=UpperCamelCase__)
             _lowerCamelCase				:  str						=		[
                 [[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(UpperCamelCase__   ,				UpperCamelCase__   ,				even_batches=UpperCamelCase__)

             _lowerCamelCase				:  List[str]						=		BatchSampler(range(20)   ,				batch_size=3   ,				drop_last=UpperCamelCase__)
             _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(UpperCamelCase__   ,				UpperCamelCase__   ,				even_batches=UpperCamelCase__)

             # Check the shards when the dataset is very small.
             _lowerCamelCase				:  Dict						=		BatchSampler(range(2)   ,				batch_size=3   ,				drop_last=UpperCamelCase__)
             _lowerCamelCase				:  Tuple						=		[[[0, 1]], []]
             self.check_batch_sampler_shards(UpperCamelCase__   ,				UpperCamelCase__   ,				even_batches=UpperCamelCase__)

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


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  Union[str, Any]						=		BatchSampler(range(24)   ,				batch_size=4   ,				drop_last=UpperCamelCase__)
             _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(UpperCamelCase__   ,				UpperCamelCase__   ,				split_batches=UpperCamelCase__   ,				even_batches=UpperCamelCase__)

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

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

             _lowerCamelCase				:  int						=		BatchSampler(range(22)   ,				batch_size=4   ,				drop_last=UpperCamelCase__)
             _lowerCamelCase				:  Tuple						=		[
                 [[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(UpperCamelCase__   ,				UpperCamelCase__   ,				split_batches=UpperCamelCase__   ,				even_batches=UpperCamelCase__)

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

             _lowerCamelCase				:  Dict						=		BatchSampler(range(21)   ,				batch_size=4   ,				drop_last=UpperCamelCase__)
             _lowerCamelCase				:  Tuple						=		[
                 [[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(UpperCamelCase__   ,				UpperCamelCase__   ,				split_batches=UpperCamelCase__   ,				even_batches=UpperCamelCase__)

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

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


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  Union[str, Any]						=		[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
             _lowerCamelCase				:  Union[str, Any]						=		[BatchSamplerShard(UpperCamelCase__   ,				2   ,				UpperCamelCase__   ,				even_batches=UpperCamelCase__) 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(UpperCamelCase__)
             _lowerCamelCase				:  Tuple						=		list(UpperCamelCase__)

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

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

             _lowerCamelCase				:  Dict						=		[]
             for idx in range(0   ,				len(UpperCamelCase__)   ,				UpperCamelCase__):
                    for l in iterable_dataset_lists:
                           observed += l[idx : idx + shard_batch_size]

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


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  List[Any]						=		42
             _lowerCamelCase				:  Any						=		RandomIterableDataset()

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

             # Edge case with a very small dataset
             _lowerCamelCase				:  str						=		RandomIterableDataset(max_length=2)

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


      def 					__snake_case      (		self):



             """simple docstring"""





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


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  Optional[Any]						=		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				:  int						=		DataLoader(list(range(16))   ,				batch_size=4)
             _lowerCamelCase				:  Union[str, Any]						=		skip_first_batches(UpperCamelCase__   ,				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				:  Any						=		DataLoaderShard(list(range(16))   ,				batch_size=4)
             for idx, _ in enumerate(UpperCamelCase__):
                    self.assertEqual(dataloader.end_of_dataloader   ,				idx == 3)

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



      def 					__snake_case      (		self):



             """simple docstring"""





             Accelerator()
             _lowerCamelCase				:  Optional[Any]						=		DataLoaderDispatcher(range(16)   ,				batch_size=4)
             for idx, _ in enumerate(UpperCamelCase__):
                    self.assertEqual(dataloader.end_of_dataloader   ,				idx == 3)

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




'''simple docstring'''







import warnings
from typing import List

import numpy as np

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding
from ...utils import is_flax_available, is_tf_available, is_torch_available


class        A__   (   _snake_case      ):
    lowercase						       =		["image_processor", "tokenizer"]
    lowercase						       =		"OwlViTImageProcessor"
    lowercase						       =		("CLIPTokenizer", "CLIPTokenizerFast")



    def __init__(       self		,					UpperCamelCase__=None		,					UpperCamelCase__=None		,					**UpperCamelCase__							)    ->   str:





           '''simple docstring'''
           A_						       =   None
           if "feature_extractor" in kwargs:
                  warnings.warn(
                      """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
                      """ instead."""		,					UpperCamelCase__		,					)
                  A_						       =   kwargs.pop("""feature_extractor"""							)

           A_						       =   image_processor if image_processor is not None else feature_extractor
           if image_processor is None:
                  raise ValueError("""You need to specify an `image_processor`."""							)
           if tokenizer is None:
                  raise ValueError("""You need to specify a `tokenizer`."""							)

           super().__init__(UpperCamelCase__		,					UpperCamelCase__							)



    def __call__(       self		,					UpperCamelCase__=None		,					UpperCamelCase__=None		,					UpperCamelCase__=None		,					UpperCamelCase__="max_length"		,					UpperCamelCase__="np"		,					**UpperCamelCase__							)    ->   int:





           '''simple docstring'''

           if text is None and query_images is None and images is None:
                  raise ValueError(
                      """You have to specify at least one text or query image or image. All three cannot be none."""							)

           if text is not None:
                  if isinstance(UpperCamelCase__		,					UpperCamelCase__							) or (isinstance(UpperCamelCase__		,					UpperCamelCase__							) and not isinstance(text[0]		,					UpperCamelCase__							)):
                         A_						       =   [self.tokenizer(UpperCamelCase__		,					padding=UpperCamelCase__		,					return_tensors=UpperCamelCase__		,					**UpperCamelCase__							)]

                  elif isinstance(UpperCamelCase__		,					UpperCamelCase__							) and isinstance(text[0]		,					UpperCamelCase__							):
                         A_						       =   []

                         # Maximum number of queries across batch
                         A_						       =   max([len(UpperCamelCase__							) for t in text]							)

                         # Pad all batch samples to max number of text queries
                         for t in text:
                                if len(UpperCamelCase__							) != max_num_queries:
                                       A_						       =   t + [""" """] * (max_num_queries - len(UpperCamelCase__							))

                                A_						       =   self.tokenizer(UpperCamelCase__		,					padding=UpperCamelCase__		,					return_tensors=UpperCamelCase__		,					**UpperCamelCase__							)
                                encodings.append(UpperCamelCase__							)
                  else:
                         raise TypeError("""Input text should be a string, a list of strings or a nested list of strings"""							)

                  if return_tensors == "np":
                         A_						       =   np.concatenate([encoding["""input_ids"""] for encoding in encodings]		,					axis=0							)
                         A_						       =   np.concatenate([encoding["""attention_mask"""] for encoding in encodings]		,					axis=0							)

                  elif return_tensors == "jax" and is_flax_available():
                         import jax.numpy as jnp

                         A_						       =   jnp.concatenate([encoding["""input_ids"""] for encoding in encodings]		,					axis=0							)
                         A_						       =   jnp.concatenate([encoding["""attention_mask"""] for encoding in encodings]		,					axis=0							)

                  elif return_tensors == "pt" and is_torch_available():
                         import torch

                         A_						       =   torch.cat([encoding["""input_ids"""] for encoding in encodings]		,					dim=0							)
                         A_						       =   torch.cat([encoding["""attention_mask"""] for encoding in encodings]		,					dim=0							)

                  elif return_tensors == "tf" and is_tf_available():
                         import tensorflow as tf

                         A_						       =   tf.stack([encoding["""input_ids"""] for encoding in encodings]		,					axis=0							)
                         A_						       =   tf.stack([encoding["""attention_mask"""] for encoding in encodings]		,					axis=0							)

                  else:
                         raise ValueError("""Target return tensor type could not be returned"""							)

                  A_						       =   BatchEncoding()
                  A_						       =   input_ids
                  A_						       =   attention_mask

           if query_images is not None:
                  A_						       =   BatchEncoding()
                  A_						       =   self.image_processor(
                      UpperCamelCase__		,					return_tensors=UpperCamelCase__		,					**UpperCamelCase__							).pixel_values
                  A_						       =   query_pixel_values

           if images is not None:
                  A_						       =   self.image_processor(UpperCamelCase__		,					return_tensors=UpperCamelCase__		,					**UpperCamelCase__							)

           if text is not None and images is not None:
                  A_						       =   image_features.pixel_values
                  return encoding
           elif query_images is not None and images is not None:
                  A_						       =   image_features.pixel_values
                  return encoding
           elif text is not None or query_images is not None:
                  return encoding
           else:
                  return BatchEncoding(data=dict(**UpperCamelCase__							)		,					tensor_type=UpperCamelCase__							)



    def        snake_case_      (       self		,					*UpperCamelCase__		,					**UpperCamelCase__							)    ->   int:





           '''simple docstring'''
           return self.image_processor.post_process(*UpperCamelCase__		,					**UpperCamelCase__							)



    def        snake_case_      (       self		,					*UpperCamelCase__		,					**UpperCamelCase__							)    ->   List[Any]:





           '''simple docstring'''
           return self.image_processor.post_process_object_detection(*UpperCamelCase__		,					**UpperCamelCase__							)



    def        snake_case_      (       self		,					*UpperCamelCase__		,					**UpperCamelCase__							)    ->   Tuple:





           '''simple docstring'''
           return self.image_processor.post_process_image_guided_detection(*UpperCamelCase__		,					**UpperCamelCase__							)



    def        snake_case_      (       self		,					*UpperCamelCase__		,					**UpperCamelCase__							)    ->   Any:





           '''simple docstring'''
           return self.tokenizer.batch_decode(*UpperCamelCase__		,					**UpperCamelCase__							)



    def        snake_case_      (       self		,					*UpperCamelCase__		,					**UpperCamelCase__							)    ->   List[str]:





           '''simple docstring'''
           return self.tokenizer.decode(*UpperCamelCase__		,					**UpperCamelCase__							)



    @property
    def        snake_case_      (       self							)    ->   Optional[int]:





           '''simple docstring'''
           warnings.warn(
               """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""		,					UpperCamelCase__		,					)
           return self.image_processor_class







    @property
    def        snake_case_      (       self							)    ->   List[str]:





           '''simple docstring'''
           warnings.warn(
               """`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."""		,					UpperCamelCase__		,					)
           return self.image_processor
