


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	:					Optional[int]		,					SCREAMING_SNAKE_CASE_	:					str=0.01		,					SCREAMING_SNAKE_CASE_	:					Any=1_0_0_0		):
													_a     			=      p_stop
													_a     			=      max_length







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




class        _UpperCamelCase			(   unittest.TestCase						):







							'''simple docstring'''







							def 		_UpperCAmelCase						(		self	:					Optional[int]		,					SCREAMING_SNAKE_CASE_	:					List[str]		,					SCREAMING_SNAKE_CASE_	:					Optional[Any]		,					SCREAMING_SNAKE_CASE_	:					Optional[int]=False		,					SCREAMING_SNAKE_CASE_	:					Optional[int]=True		):
													_a     			=      [
													    BatchSamplerShard(lowerCamelCase_		,					2		,					lowerCamelCase_		,					split_batches=lowerCamelCase_		,					even_batches=lowerCamelCase_		)
													    for i in range(2		)
													]
													_a     			=      [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 		_UpperCAmelCase						(		self	:					List[Any]		):
													# Check the shards when the dataset is a round multiple of total batch size.
													_a     			=      BatchSampler(range(2_4		)		,					batch_size=3		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
													    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1, 2_2, 2_3]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		)

													_a     			=      BatchSampler(range(2_4		)		,					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.
													_a     			=      BatchSampler(range(2_1		)		,					batch_size=3		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
													    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [0, 1, 2]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		)

													_a     			=      BatchSampler(range(2_1		)		,					batch_size=3		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
													    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
													]
													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.
													_a     			=      BatchSampler(range(2_2		)		,					batch_size=3		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
													    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1, 0, 1]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		)

													_a     			=      BatchSampler(range(2_2		)		,					batch_size=3		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
													    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
													]
													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.
													_a     			=      BatchSampler(range(2_0		)		,					batch_size=3		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 0]],
													    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [1, 2, 3]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		)

													_a     			=      BatchSampler(range(2_0		)		,					batch_size=3		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
													    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		)

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

													_a     			=      BatchSampler(range(2		)		,					batch_size=3		,					drop_last=lowerCamelCase_		)
													_a     			=      [[], []]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		)







							def 		_UpperCAmelCase						(		self	:					Union[str, Any]		):
													# Check the shards when the dataset is a round multiple of batch size.
													_a     			=      BatchSampler(range(2_4		)		,					batch_size=4		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
													    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [2_2, 2_3]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		,					split_batches=lowerCamelCase_		)

													_a     			=      BatchSampler(range(2_4		)		,					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.
													_a     			=      BatchSampler(range(2_2		)		,					batch_size=4		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
													    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [0, 1]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		,					split_batches=lowerCamelCase_		)

													_a     			=      BatchSampler(range(2_2		)		,					batch_size=4		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
													    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
													]
													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.
													_a     			=      BatchSampler(range(2_1		)		,					batch_size=4		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 0]],
													    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [1, 2]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		,					split_batches=lowerCamelCase_		)

													_a     			=      BatchSampler(range(2_1		)		,					batch_size=4		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
													    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		,					split_batches=lowerCamelCase_		)

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

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







							def 		_UpperCAmelCase						(		self	:					List[str]		):
													# Check the shards when the dataset is a round multiple of total batch size.
													_a     			=      BatchSampler(range(2_4		)		,					batch_size=3		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
													    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1, 2_2, 2_3]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		,					even_batches=lowerCamelCase_		)

													_a     			=      BatchSampler(range(2_4		)		,					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.
													_a     			=      BatchSampler(range(2_1		)		,					batch_size=3		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
													    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		,					even_batches=lowerCamelCase_		)

													_a     			=      BatchSampler(range(2_1		)		,					batch_size=3		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
													    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
													]
													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.
													_a     			=      BatchSampler(range(2_2		)		,					batch_size=3		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
													    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		,					even_batches=lowerCamelCase_		)

													_a     			=      BatchSampler(range(2_2		)		,					batch_size=3		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
													    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
													]
													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.
													_a     			=      BatchSampler(range(2_0		)		,					batch_size=3		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9]],
													    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		,					even_batches=lowerCamelCase_		)

													_a     			=      BatchSampler(range(2_0		)		,					batch_size=3		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
													    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		,					even_batches=lowerCamelCase_		)

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

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







							def 		_UpperCAmelCase						(		self	:					Union[str, Any]		):
													# Check the shards when the dataset is a round multiple of batch size.
													_a     			=      BatchSampler(range(2_4		)		,					batch_size=4		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
													    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [2_2, 2_3]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		,					split_batches=lowerCamelCase_		,					even_batches=lowerCamelCase_		)

													_a     			=      BatchSampler(range(2_4		)		,					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.
													_a     			=      BatchSampler(range(2_2		)		,					batch_size=4		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
													    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		,					split_batches=lowerCamelCase_		,					even_batches=lowerCamelCase_		)

													_a     			=      BatchSampler(range(2_2		)		,					batch_size=4		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
													    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
													]
													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.
													_a     			=      BatchSampler(range(2_1		)		,					batch_size=4		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0]],
													    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		,					split_batches=lowerCamelCase_		,					even_batches=lowerCamelCase_		)

													_a     			=      BatchSampler(range(2_1		)		,					batch_size=4		,					drop_last=lowerCamelCase_		)
													_a     			=      [
													    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
													    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
													]
													self.check_batch_sampler_shards(lowerCamelCase_		,					lowerCamelCase_		,					split_batches=lowerCamelCase_		,					even_batches=lowerCamelCase_		)

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

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







							def 		_UpperCAmelCase						(		self	:					Any		):
													_a     			=      [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 1_0, 1_1], [1_2, 1_3]]
													_a     			=      [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], [1_2, 1_3]]		)
													self.assertListEqual(list(batch_sampler_shards[1]		)		,					[[3, 4], [9, 1_0, 1_1]]		)







							def 		_UpperCAmelCase						(		self	:					Dict		,					SCREAMING_SNAKE_CASE_	:					List[Any]		,					SCREAMING_SNAKE_CASE_	:					Optional[int]		,					SCREAMING_SNAKE_CASE_	:					str		,					SCREAMING_SNAKE_CASE_	:					Any=False		,					SCREAMING_SNAKE_CASE_	:					Union[str, Any]=2		,					SCREAMING_SNAKE_CASE_	:					Tuple=False		):
													random.seed(lowerCamelCase_		)
													_a     			=      list(lowerCamelCase_		)

													_a     			=      [
													    IterableDatasetShard(
													        lowerCamelCase_		,					batch_size=lowerCamelCase_		,					drop_last=lowerCamelCase_		,					num_processes=lowerCamelCase_		,					process_index=lowerCamelCase_		,					split_batches=lowerCamelCase_		,					)
													    for i in range(lowerCamelCase_		)
													]
													_a     			=      []
													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_		)		)

													_a     			=      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
													_a     			=      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		)

													_a     			=      []
													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 		_UpperCAmelCase						(		self	:					str		):
													_a     			=      4_2
													_a     			=      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
													_a     			=      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 		_UpperCAmelCase						(		self	:					int		):
													_a     			=      BatchSampler(range(1_6		)		,					batch_size=4		,					drop_last=lowerCamelCase_		)
													_a     			=      SkipBatchSampler(lowerCamelCase_		,					2		)
													self.assertListEqual(list(lowerCamelCase_		)		,					[[8, 9, 1_0, 1_1], [1_2, 1_3, 1_4, 1_5]]		)







							def 		_UpperCAmelCase						(		self	:					Union[str, Any]		):
													_a     			=      SkipDataLoader(list(range(1_6		)		)		,					batch_size=4		,					skip_batches=2		)
													self.assertListEqual([t.tolist() for t in dataloader]		,					[[8, 9, 1_0, 1_1], [1_2, 1_3, 1_4, 1_5]]		)







							def 		_UpperCAmelCase						(		self	:					Dict		):
													_a     			=      DataLoader(list(range(1_6		)		)		,					batch_size=4		)
													_a     			=      skip_first_batches(lowerCamelCase_		,					num_batches=2		)
													self.assertListEqual([t.tolist() for t in new_dataloader]		,					[[8, 9, 1_0, 1_1], [1_2, 1_3, 1_4, 1_5]]		)







							def 		_UpperCAmelCase						(		self	:					Optional[int]		):
													_a     			=      DataLoaderShard(list(range(1_6		)		)		,					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 		_UpperCAmelCase						(		self	:					int		):
													Accelerator()
													_a     			=      DataLoaderDispatcher(range(1_6		)		,					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		)






'''simple docstring'''






import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, XLMRobertaTokenizer

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class 						a__							(				a__							,							a__							,							a__							,							unittest.TestCase    ):





							'''simple docstring'''






							lowercase__     :  Optional[Any]      		= AltDiffusionPipeline
							lowercase__     :  Dict      		= TEXT_TO_IMAGE_PARAMS
							lowercase__     :  str      		= TEXT_TO_IMAGE_BATCH_PARAMS
							lowercase__     :  Any      		= TEXT_TO_IMAGE_IMAGE_PARAMS
							lowercase__     :  List[str]      		= TEXT_TO_IMAGE_IMAGE_PARAMS







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														torch.manual_seed(0  )
														lowerCAmelCase__			  =      UNetaDConditionModel(
														    block_out_channels=(32, 64)   ,       layers_per_block=2   ,       sample_size=32   ,       in_channels=4   ,       out_channels=4   ,       down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')   ,       up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')   ,       cross_attention_dim=32   ,       )
														lowerCAmelCase__			  =      DDIMScheduler(
														    beta_start=0.00_085   ,       beta_end=0.012   ,       beta_schedule='''scaled_linear'''   ,       clip_sample=lowerCamelCase_   ,       set_alpha_to_one=lowerCamelCase_   ,       )
														torch.manual_seed(0  )
														lowerCAmelCase__			  =      AutoencoderKL(
														    block_out_channels=[32, 64]   ,       in_channels=3   ,       out_channels=3   ,       down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,       up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,       latent_channels=4   ,       )

														# TODO: address the non-deterministic text encoder (fails for save-load tests)
														# torch.manual_seed(0)
														# text_encoder_config = RobertaSeriesConfig(
														#     hidden_size=32,
														#     project_dim=32,
														#     intermediate_size=37,
														#     layer_norm_eps=1e-05,
														#     num_attention_heads=4,
														#     num_hidden_layers=5,
														#     vocab_size=5002,
														# )
														# text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

														torch.manual_seed(0  )
														lowerCAmelCase__			  =      CLIPTextConfig(
														    bos_token_id=0   ,       eos_token_id=2   ,       hidden_size=32   ,       projection_dim=32   ,       intermediate_size=37   ,       layer_norm_eps=1e-05   ,       num_attention_heads=4   ,       num_hidden_layers=5   ,       pad_token_id=1   ,       vocab_size=50_02   ,       )
														lowerCAmelCase__			  =      CLIPTextModel(lowerCamelCase_  )

														lowerCAmelCase__			  =      XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''  )
														lowerCAmelCase__			  =      77

														lowerCAmelCase__			  =      {
														    '''unet''': unet,
														    '''scheduler''': scheduler,
														    '''vae''': vae,
														    '''text_encoder''': text_encoder,
														    '''tokenizer''': tokenizer,
														    '''safety_checker''': None,
														    '''feature_extractor''': None,
														}
														return components







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_=0  )			->				List[str]:
														if str(lowerCamelCase_  ).startswith('''mps'''  ):
																					lowerCAmelCase__			  =      torch.manual_seed(lowerCamelCase_  )
														else:
																					lowerCAmelCase__			  =      torch.Generator(device=lowerCamelCase_  ).manual_seed(lowerCamelCase_  )
														lowerCAmelCase__			  =      {
														    '''prompt''': '''A painting of a squirrel eating a burger''',
														    '''generator''': generator,
														    '''num_inference_steps''': 2,
														    '''guidance_scale''': 6.0,
														    '''output_type''': '''numpy''',
														}
														return inputs







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Union[str, Any]:
														super().test_attention_slicing_forward_pass(expected_max_diff=3e-3  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[int]:
														super().test_inference_batch_single_identical(expected_max_diff=3e-3  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[str]:
														lowerCAmelCase__			  =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator

														lowerCAmelCase__			  =      self.get_dummy_components()
														torch.manual_seed(0  )
														lowerCAmelCase__			  =      RobertaSeriesConfig(
														    hidden_size=32   ,       project_dim=32   ,       intermediate_size=37   ,       layer_norm_eps=1e-05   ,       num_attention_heads=4   ,       num_hidden_layers=5   ,       vocab_size=50_02   ,       )
														# TODO: remove after fixing the non-deterministic text encoder
														lowerCAmelCase__			  =      RobertaSeriesModelWithTransformation(lowerCamelCase_  )
														lowerCAmelCase__			  =      text_encoder

														lowerCAmelCase__			  =      AltDiffusionPipeline(**lowerCamelCase_  )
														lowerCAmelCase__			  =      alt_pipe.to(lowerCamelCase_  )
														alt_pipe.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      self.get_dummy_inputs(lowerCamelCase_  )
														lowerCAmelCase__			  =      '''A photo of an astronaut'''
														lowerCAmelCase__			  =      alt_pipe(**lowerCamelCase_  )
														lowerCAmelCase__			  =      output.images
														lowerCAmelCase__			  =      image[0, -3:, -3:, -1]

														assert image.shape == (1, 64, 64, 3)
														lowerCAmelCase__			  =      np.array(
														    [0.5_748_162, 0.60_447_145, 0.48_821_217, 0.50_100_636, 0.5_431_185, 0.45_763_683, 0.49_657_696, 0.48_132_733, 0.47_573_093]  )

														assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														lowerCAmelCase__			  =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator

														lowerCAmelCase__			  =      self.get_dummy_components()
														lowerCAmelCase__			  =      PNDMScheduler(skip_prk_steps=lowerCamelCase_  )
														torch.manual_seed(0  )
														lowerCAmelCase__			  =      RobertaSeriesConfig(
														    hidden_size=32   ,       project_dim=32   ,       intermediate_size=37   ,       layer_norm_eps=1e-05   ,       num_attention_heads=4   ,       num_hidden_layers=5   ,       vocab_size=50_02   ,       )
														# TODO: remove after fixing the non-deterministic text encoder
														lowerCAmelCase__			  =      RobertaSeriesModelWithTransformation(lowerCamelCase_  )
														lowerCAmelCase__			  =      text_encoder
														lowerCAmelCase__			  =      AltDiffusionPipeline(**lowerCamelCase_  )
														lowerCAmelCase__			  =      alt_pipe.to(lowerCamelCase_  )
														alt_pipe.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      self.get_dummy_inputs(lowerCamelCase_  )
														lowerCAmelCase__			  =      alt_pipe(**lowerCamelCase_  )
														lowerCAmelCase__			  =      output.images
														lowerCAmelCase__			  =      image[0, -3:, -3:, -1]

														assert image.shape == (1, 64, 64, 3)
														lowerCAmelCase__			  =      np.array(
														    [0.51_605_093, 0.5_707_241, 0.47_365_507, 0.50_578_886, 0.5_633_877, 0.4_642_503, 0.5_182_081, 0.48_763_484, 0.49_084_237]  )

														assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2




@slow
@require_torch_gpu
class 						a__							(				unittest.TestCase    ):





							'''simple docstring'''






							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Tuple:
														# clean up the VRAM after each test
														super().tearDown()
														gc.collect()
														torch.cuda.empty_cache()







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														# make sure here that pndm scheduler skips prk
														lowerCAmelCase__			  =      AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''   ,       safety_checker=lowerCamelCase_  )
														lowerCAmelCase__			  =      alt_pipe.to(lowerCamelCase_  )
														alt_pipe.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      '''A painting of a squirrel eating a burger'''
														lowerCAmelCase__			  =      torch.manual_seed(0  )
														lowerCAmelCase__			  =      alt_pipe([prompt]   ,       generator=lowerCamelCase_   ,       guidance_scale=6.0   ,       num_inference_steps=20   ,       output_type='''np'''  )

														lowerCAmelCase__			  =      output.images

														lowerCAmelCase__			  =      image[0, -3:, -3:, -1]

														assert image.shape == (1, 5_12, 5_12, 3)
														lowerCAmelCase__			  =      np.array([0.1_010, 0.0_800, 0.0_794, 0.0_885, 0.0_843, 0.0_762, 0.0_769, 0.0_729, 0.0_586]  )

														assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														lowerCAmelCase__			  =      DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''   ,       subfolder='''scheduler'''  )

														lowerCAmelCase__			  =      AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''   ,       scheduler=lowerCamelCase_   ,       safety_checker=lowerCamelCase_  )
														lowerCAmelCase__			  =      alt_pipe.to(lowerCamelCase_  )
														alt_pipe.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      '''A painting of a squirrel eating a burger'''
														lowerCAmelCase__			  =      torch.manual_seed(0  )

														lowerCAmelCase__			  =      alt_pipe([prompt]   ,       generator=lowerCamelCase_   ,       num_inference_steps=2   ,       output_type='''numpy'''  )
														lowerCAmelCase__			  =      output.images

														lowerCAmelCase__			  =      image[0, -3:, -3:, -1]

														assert image.shape == (1, 5_12, 5_12, 3)
														lowerCAmelCase__			  =      np.array([0.4_019, 0.4_052, 0.3_810, 0.4_119, 0.3_916, 0.3_982, 0.4_651, 0.4_195, 0.5_323]  )

														assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2