



'''simple docstring'''
import tempfile

import torch

from diffusers import (
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    UniPCMultistepScheduler,
)

from .test_schedulers import SchedulerCommonTest






class 							lowercase__	(							lowercase   ):
	lowercase__  				=			(DEISMultistepScheduler,)
	lowercase__  				=			(("""num_inference_steps""", 25),)





	def 	UpperCamelCase_  ( self					:				Dict ,**lowerCamelCase__					:				Optional[int] ):







							'''simple docstring'''
							_UpperCamelCase					:		Union[str, Any] 				=   {
							    'num_train_timesteps': 1000,
							    'beta_start': 0.0_0_0_1,
							    'beta_end': 0.0_2,
							    'beta_schedule': 'linear',
							    'solver_order': 2,
							}

							config.update(**lowerCamelCase__ )
							return config





	def 	UpperCamelCase_  ( self					:				List[Any] ,lowerCamelCase__					:				Tuple=0 ,**lowerCamelCase__					:				List[str] ):







							'''simple docstring'''
							_UpperCamelCase					:		Dict 				=   dict(self.forward_default_kwargs )
							_UpperCamelCase					:		int 				=   kwargs.pop('num_inference_steps' ,lowerCamelCase__ )
							_UpperCamelCase					:		int 				=   self.dummy_sample
							_UpperCamelCase					:		Union[str, Any] 				=   0.1 * sample
							_UpperCamelCase					:		str 				=   [residual + 0.2, residual + 0.1_5, residual + 0.1_0]

							for scheduler_class in self.scheduler_classes:
													_UpperCamelCase					:		Optional[Any] 				=   self.get_scheduler_config(**lowerCamelCase__ )
													_UpperCamelCase					:		Optional[int] 				=   scheduler_class(**lowerCamelCase__ )
													scheduler.set_timesteps(lowerCamelCase__ )
													# copy over dummy past residuals
													_UpperCamelCase					:		Optional[Any] 				=   dummy_past_residuals[: scheduler.config.solver_order]

													with tempfile.TemporaryDirectory() as tmpdirname:
																			scheduler.save_config(lowerCamelCase__ )
																			_UpperCamelCase					:		Tuple 				=   scheduler_class.from_pretrained(lowerCamelCase__ )
																			new_scheduler.set_timesteps(lowerCamelCase__ )
																			# copy over dummy past residuals
																			_UpperCamelCase					:		str 				=   dummy_past_residuals[: new_scheduler.config.solver_order]

													_UpperCamelCase				,			_UpperCamelCase					:		Any 				=   sample, sample
													for t in range(lowerCamelCase__ ,time_step + scheduler.config.solver_order + 1 ):
																			_UpperCamelCase					:		int 				=   scheduler.step(lowerCamelCase__ ,lowerCamelCase__ ,lowerCamelCase__ ,**lowerCamelCase__ ).prev_sample
																			_UpperCamelCase					:		Optional[Any] 				=   new_scheduler.step(lowerCamelCase__ ,lowerCamelCase__ ,lowerCamelCase__ ,**lowerCamelCase__ ).prev_sample

																			assert torch.sum(torch.abs(output - new_output ) ) < 1E-5, "Scheduler outputs are not identical"





	def 	UpperCamelCase_  ( self					:				Tuple ):







							'''simple docstring'''
							pass





	def 	UpperCamelCase_  ( self					:				List[Any] ,lowerCamelCase__					:				List[str]=0 ,**lowerCamelCase__					:				Dict ):







							'''simple docstring'''
							_UpperCamelCase					:		List[str] 				=   dict(self.forward_default_kwargs )
							_UpperCamelCase					:		List[Any] 				=   kwargs.pop('num_inference_steps' ,lowerCamelCase__ )
							_UpperCamelCase					:		int 				=   self.dummy_sample
							_UpperCamelCase					:		str 				=   0.1 * sample
							_UpperCamelCase					:		Tuple 				=   [residual + 0.2, residual + 0.1_5, residual + 0.1_0]

							for scheduler_class in self.scheduler_classes:
													_UpperCamelCase					:		Optional[int] 				=   self.get_scheduler_config()
													_UpperCamelCase					:		Tuple 				=   scheduler_class(**lowerCamelCase__ )
													scheduler.set_timesteps(lowerCamelCase__ )

													# copy over dummy past residuals (must be after setting timesteps)
													_UpperCamelCase					:		List[str] 				=   dummy_past_residuals[: scheduler.config.solver_order]

													with tempfile.TemporaryDirectory() as tmpdirname:
																			scheduler.save_config(lowerCamelCase__ )
																			_UpperCamelCase					:		Optional[int] 				=   scheduler_class.from_pretrained(lowerCamelCase__ )
																			# copy over dummy past residuals
																			new_scheduler.set_timesteps(lowerCamelCase__ )

																			# copy over dummy past residual (must be after setting timesteps)
																			_UpperCamelCase					:		Optional[Any] 				=   dummy_past_residuals[: new_scheduler.config.solver_order]

													_UpperCamelCase					:		Optional[int] 				=   scheduler.step(lowerCamelCase__ ,lowerCamelCase__ ,lowerCamelCase__ ,**lowerCamelCase__ ).prev_sample
													_UpperCamelCase					:		List[Any] 				=   new_scheduler.step(lowerCamelCase__ ,lowerCamelCase__ ,lowerCamelCase__ ,**lowerCamelCase__ ).prev_sample

													assert torch.sum(torch.abs(output - new_output ) ) < 1E-5, "Scheduler outputs are not identical"





	def 	UpperCamelCase_  ( self					:				Any ,lowerCamelCase__					:				int=None ,**lowerCamelCase__					:				Any ):







							'''simple docstring'''
							if scheduler is None:
													_UpperCamelCase					:		Any 				=   self.scheduler_classes[0]
													_UpperCamelCase					:		Optional[Any] 				=   self.get_scheduler_config(**lowerCamelCase__ )
													_UpperCamelCase					:		str 				=   scheduler_class(**lowerCamelCase__ )

							_UpperCamelCase					:		List[Any] 				=   self.scheduler_classes[0]
							_UpperCamelCase					:		str 				=   self.get_scheduler_config(**lowerCamelCase__ )
							_UpperCamelCase					:		int 				=   scheduler_class(**lowerCamelCase__ )

							_UpperCamelCase					:		Dict 				=   10
							_UpperCamelCase					:		Optional[int] 				=   self.dummy_model()
							_UpperCamelCase					:		Tuple 				=   self.dummy_sample_deter
							scheduler.set_timesteps(lowerCamelCase__ )

							for i, t in enumerate(scheduler.timesteps ):
													_UpperCamelCase					:		Dict 				=   model(lowerCamelCase__ ,lowerCamelCase__ )
													_UpperCamelCase					:		List[str] 				=   scheduler.step(lowerCamelCase__ ,lowerCamelCase__ ,lowerCamelCase__ ).prev_sample

							return sample





	def 	UpperCamelCase_  ( self					:				Optional[int] ):







							'''simple docstring'''
							_UpperCamelCase					:		Any 				=   dict(self.forward_default_kwargs )

							_UpperCamelCase					:		Dict 				=   kwargs.pop('num_inference_steps' ,lowerCamelCase__ )

							for scheduler_class in self.scheduler_classes:
													_UpperCamelCase					:		Union[str, Any] 				=   self.get_scheduler_config()
													_UpperCamelCase					:		Optional[int] 				=   scheduler_class(**lowerCamelCase__ )

													_UpperCamelCase					:		str 				=   self.dummy_sample
													_UpperCamelCase					:		Optional[Any] 				=   0.1 * sample

													if num_inference_steps is not None and hasattr(lowerCamelCase__ ,'set_timesteps' ):
																			scheduler.set_timesteps(lowerCamelCase__ )
													elif num_inference_steps is not None and not hasattr(lowerCamelCase__ ,'set_timesteps' ):
																			_UpperCamelCase					:		List[Any] 				=   num_inference_steps

													# copy over dummy past residuals (must be done after set_timesteps)
													_UpperCamelCase					:		Tuple 				=   [residual + 0.2, residual + 0.1_5, residual + 0.1_0]
													_UpperCamelCase					:		List[str] 				=   dummy_past_residuals[: scheduler.config.solver_order]

													_UpperCamelCase					:		Union[str, Any] 				=   scheduler.timesteps[5]
													_UpperCamelCase					:		int 				=   scheduler.timesteps[6]

													_UpperCamelCase					:		Tuple 				=   scheduler.step(lowerCamelCase__ ,lowerCamelCase__ ,lowerCamelCase__ ,**lowerCamelCase__ ).prev_sample
													_UpperCamelCase					:		str 				=   scheduler.step(lowerCamelCase__ ,lowerCamelCase__ ,lowerCamelCase__ ,**lowerCamelCase__ ).prev_sample

													self.assertEqual(output_a.shape ,sample.shape )
													self.assertEqual(output_a.shape ,output_a.shape )





	def 	UpperCamelCase_  ( self					:				Union[str, Any] ):







							'''simple docstring'''
							# make sure that iterating over schedulers with same config names gives same results
							# for defaults
							_UpperCamelCase					:		List[str] 				=   DEISMultistepScheduler(**self.get_scheduler_config() )
							_UpperCamelCase					:		Any 				=   self.full_loop(scheduler=lowerCamelCase__ )
							_UpperCamelCase					:		List[Any] 				=   torch.mean(torch.abs(lowerCamelCase__ ) )

							assert abs(result_mean.item() - 0.2_3_9_1_6 ) < 1E-3

							_UpperCamelCase					:		str 				=   DPMSolverSinglestepScheduler.from_config(scheduler.config )
							_UpperCamelCase					:		List[str] 				=   DPMSolverMultistepScheduler.from_config(scheduler.config )
							_UpperCamelCase					:		Union[str, Any] 				=   UniPCMultistepScheduler.from_config(scheduler.config )
							_UpperCamelCase					:		List[str] 				=   DEISMultistepScheduler.from_config(scheduler.config )

							_UpperCamelCase					:		Dict 				=   self.full_loop(scheduler=lowerCamelCase__ )
							_UpperCamelCase					:		Union[str, Any] 				=   torch.mean(torch.abs(lowerCamelCase__ ) )

							assert abs(result_mean.item() - 0.2_3_9_1_6 ) < 1E-3





	def 	UpperCamelCase_  ( self					:				Optional[int] ):







							'''simple docstring'''
							for timesteps in [25, 50, 100, 999, 1000]:
													self.check_over_configs(num_train_timesteps=lowerCamelCase__ )





	def 	UpperCamelCase_  ( self					:				Any ):







							'''simple docstring'''
							self.check_over_configs(thresholding=lowerCamelCase__ )
							for order in [1, 2, 3]:
													for solver_type in ["logrho"]:
																			for threshold in [0.5, 1.0, 2.0]:
																									for prediction_type in ["epsilon", "sample"]:
																															self.check_over_configs(
																															    thresholding=lowerCamelCase__ ,prediction_type=lowerCamelCase__ ,sample_max_value=lowerCamelCase__ ,algorithm_type='deis' ,solver_order=lowerCamelCase__ ,solver_type=lowerCamelCase__ ,)





	def 	UpperCamelCase_  ( self					:				List[str] ):







							'''simple docstring'''
							for prediction_type in ["epsilon", "v_prediction"]:
													self.check_over_configs(prediction_type=lowerCamelCase__ )





	def 	UpperCamelCase_  ( self					:				Dict ):







							'''simple docstring'''
							for algorithm_type in ["deis"]:
													for solver_type in ["logrho"]:
																			for order in [1, 2, 3]:
																									for prediction_type in ["epsilon", "sample"]:
																															self.check_over_configs(
																															    solver_order=lowerCamelCase__ ,solver_type=lowerCamelCase__ ,prediction_type=lowerCamelCase__ ,algorithm_type=lowerCamelCase__ ,)
																															_UpperCamelCase					:		int 				=   self.full_loop(
																															    solver_order=lowerCamelCase__ ,solver_type=lowerCamelCase__ ,prediction_type=lowerCamelCase__ ,algorithm_type=lowerCamelCase__ ,)
																															assert not torch.isnan(lowerCamelCase__ ).any(), "Samples have nan numbers"





	def 	UpperCamelCase_  ( self					:				int ):







							'''simple docstring'''
							self.check_over_configs(lower_order_final=lowerCamelCase__ )
							self.check_over_configs(lower_order_final=lowerCamelCase__ )





	def 	UpperCamelCase_  ( self					:				List[Any] ):







							'''simple docstring'''
							for num_inference_steps in [1, 2, 3, 5, 10, 50, 100, 999, 1000]:
													self.check_over_forward(num_inference_steps=lowerCamelCase__ ,time_step=0 )





	def 	UpperCamelCase_  ( self					:				Optional[Any] ):







							'''simple docstring'''
							_UpperCamelCase					:		str 				=   self.full_loop()
							_UpperCamelCase					:		Tuple 				=   torch.mean(torch.abs(lowerCamelCase__ ) )

							assert abs(result_mean.item() - 0.2_3_9_1_6 ) < 1E-3





	def 	UpperCamelCase_  ( self					:				int ):







							'''simple docstring'''
							_UpperCamelCase					:		Union[str, Any] 				=   self.full_loop(prediction_type='v_prediction' )
							_UpperCamelCase					:		Tuple 				=   torch.mean(torch.abs(lowerCamelCase__ ) )

							assert abs(result_mean.item() - 0.0_9_1 ) < 1E-3





	def 	UpperCamelCase_  ( self					:				List[str] ):







							'''simple docstring'''
							_UpperCamelCase					:		Optional[int] 				=   self.scheduler_classes[0]
							_UpperCamelCase					:		List[Any] 				=   self.get_scheduler_config(thresholding=lowerCamelCase__ ,dynamic_thresholding_ratio=0 )
							_UpperCamelCase					:		Any 				=   scheduler_class(**lowerCamelCase__ )

							_UpperCamelCase					:		Dict 				=   10
							_UpperCamelCase					:		List[Any] 				=   self.dummy_model()
							_UpperCamelCase					:		Any 				=   self.dummy_sample_deter.half()
							scheduler.set_timesteps(lowerCamelCase__ )

							for i, t in enumerate(scheduler.timesteps ):
													_UpperCamelCase					:		Optional[Any] 				=   model(lowerCamelCase__ ,lowerCamelCase__ )
													_UpperCamelCase					:		List[Any] 				=   scheduler.step(lowerCamelCase__ ,lowerCamelCase__ ,lowerCamelCase__ ).prev_sample

							assert sample.dtype == torch.floataa






"""simple docstring"""







import random
import unittest

from torch.utils.data import BatchSampler, DataLoader, IterableDataset

from accelerate import Accelerator
from accelerate.data_loader import (
    BatchSamplerShard,
    DataLoaderDispatcher,
    DataLoaderShard,
    IterableDatasetShard,
    SkipBatchSampler,
    SkipDataLoader,
    skip_first_batches,
)


class       SCREAMING_SNAKE_CASE__				(		lowercase				):
       """simple docstring"""


       def __init__(						self		,     snake_case__=0.01		,     snake_case__=1_000						):






         """simple docstring"""





         lowerCAmelCase     :     List[Any]								=							p_stop
         lowerCAmelCase     :     Optional[Any]								=							max_length







       def __iter__(						self						):






         """simple docstring"""





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


class       SCREAMING_SNAKE_CASE__				(		unittest.TestCase				):
       """simple docstring"""


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






         """simple docstring"""





         lowerCAmelCase     :     Dict								=							[
             BatchSamplerShard(snake_case__		,     2		,     snake_case__		,     split_batches=snake_case__		,     even_batches=snake_case__						)
             for i in range(2						)
         ]
         lowerCAmelCase     :     Any								=							[list(snake_case__						) for batch_sampler_shard in batch_sampler_shards]
         if not split_batches:
           self.assertListEqual([len(snake_case__						) for shard in batch_sampler_shards]		,     [len(snake_case__						) for e in expected]						)
         self.assertListEqual(snake_case__		,     snake_case__						)

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     Union[str, Any]								=							BatchSampler(range(24						)		,     batch_size=3		,     drop_last=snake_case__						)
         lowerCAmelCase     :     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(snake_case__		,     snake_case__						)

         lowerCAmelCase     :     Tuple								=							BatchSampler(range(24						)		,     batch_size=3		,     drop_last=snake_case__						)
         # Expected shouldn't change
         self.check_batch_sampler_shards(snake_case__		,     snake_case__						)

         # 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=snake_case__						)
         lowerCAmelCase     :     List[str]								=							[
             [[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(snake_case__		,     snake_case__						)

         lowerCAmelCase     :     Tuple								=							BatchSampler(range(21						)		,     batch_size=3		,     drop_last=snake_case__						)
         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(snake_case__		,     snake_case__						)

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

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

         # 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     :     Any								=							BatchSampler(range(20						)		,     batch_size=3		,     drop_last=snake_case__						)
         lowerCAmelCase     :     Tuple								=							[
             [[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(snake_case__		,     snake_case__						)

         lowerCAmelCase     :     List[str]								=							BatchSampler(range(20						)		,     batch_size=3		,     drop_last=snake_case__						)
         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(snake_case__		,     snake_case__						)

         # Check the shards when the dataset is very small.
         lowerCAmelCase     :     Dict								=							BatchSampler(range(2						)		,     batch_size=3		,     drop_last=snake_case__						)
         lowerCAmelCase     :     Union[str, Any]								=							[[[0, 1, 0]], [[1, 0, 1]]]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__						)

         lowerCAmelCase     :     Optional[Any]								=							BatchSampler(range(2						)		,     batch_size=3		,     drop_last=snake_case__						)
         lowerCAmelCase     :     Any								=							[[], []]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__						)

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     Tuple								=							BatchSampler(range(24						)		,     batch_size=4		,     drop_last=snake_case__						)
         lowerCAmelCase     :     int								=							[
             [[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(snake_case__		,     snake_case__		,     split_batches=snake_case__						)

         lowerCAmelCase     :     Dict								=							BatchSampler(range(24						)		,     batch_size=4		,     drop_last=snake_case__						)
         # Expected shouldn't change
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     split_batches=snake_case__						)

         # Check the shards when the dataset is not a round multiple of batch size.
         lowerCAmelCase     :     Optional[int]								=							BatchSampler(range(22						)		,     batch_size=4		,     drop_last=snake_case__						)
         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], [0, 1]],
         ]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     split_batches=snake_case__						)

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

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

         lowerCAmelCase     :     Any								=							BatchSampler(range(21						)		,     batch_size=4		,     drop_last=snake_case__						)
         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(snake_case__		,     snake_case__		,     split_batches=snake_case__						)

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

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

       def        lowercase__							(						self						):






         """simple docstring"""





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

         lowerCAmelCase     :     Optional[Any]								=							BatchSampler(range(24						)		,     batch_size=3		,     drop_last=snake_case__						)
         # Expected shouldn't change
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     even_batches=snake_case__						)

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

         lowerCAmelCase     :     int								=							BatchSampler(range(21						)		,     batch_size=3		,     drop_last=snake_case__						)
         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(snake_case__		,     snake_case__		,     even_batches=snake_case__						)

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

         lowerCAmelCase     :     List[str]								=							BatchSampler(range(22						)		,     batch_size=3		,     drop_last=snake_case__						)
         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(snake_case__		,     snake_case__		,     even_batches=snake_case__						)

         # 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     :     str								=							BatchSampler(range(20						)		,     batch_size=3		,     drop_last=snake_case__						)
         lowerCAmelCase     :     Union[str, Any]								=							[
             [[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(snake_case__		,     snake_case__		,     even_batches=snake_case__						)

         lowerCAmelCase     :     Any								=							BatchSampler(range(20						)		,     batch_size=3		,     drop_last=snake_case__						)
         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(snake_case__		,     snake_case__		,     even_batches=snake_case__						)

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

         lowerCAmelCase     :     str								=							BatchSampler(range(2						)		,     batch_size=3		,     drop_last=snake_case__						)
         lowerCAmelCase     :     Optional[Any]								=							[[], []]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     even_batches=snake_case__						)

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     str								=							BatchSampler(range(24						)		,     batch_size=4		,     drop_last=snake_case__						)
         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(snake_case__		,     snake_case__		,     split_batches=snake_case__		,     even_batches=snake_case__						)

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

         # Check the shards when the dataset is not a round multiple of batch size.
         lowerCAmelCase     :     Tuple								=							BatchSampler(range(22						)		,     batch_size=4		,     drop_last=snake_case__						)
         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]],
         ]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     split_batches=snake_case__		,     even_batches=snake_case__						)

         lowerCAmelCase     :     int								=							BatchSampler(range(22						)		,     batch_size=4		,     drop_last=snake_case__						)
         lowerCAmelCase     :     str								=							[
             [[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(snake_case__		,     snake_case__		,     split_batches=snake_case__		,     even_batches=snake_case__						)

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

         lowerCAmelCase     :     str								=							BatchSampler(range(21						)		,     batch_size=4		,     drop_last=snake_case__						)
         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(snake_case__		,     snake_case__		,     split_batches=snake_case__		,     even_batches=snake_case__						)

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

         lowerCAmelCase     :     Dict								=							BatchSampler(range(2						)		,     batch_size=4		,     drop_last=snake_case__						)
         lowerCAmelCase     :     List[str]								=							[[], []]
         self.check_batch_sampler_shards(snake_case__		,     snake_case__		,     split_batches=snake_case__		,     even_batches=snake_case__						)

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     List[str]								=							[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
         lowerCAmelCase     :     Tuple								=							[BatchSamplerShard(snake_case__		,     2		,     snake_case__		,     even_batches=snake_case__						) for i in range(2						)]

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

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

       def        lowercase__							(						self		,     snake_case__		,     snake_case__		,     snake_case__		,     snake_case__=False		,     snake_case__=2		,     snake_case__=False						):






         """simple docstring"""





         random.seed(snake_case__						)
         lowerCAmelCase     :     List[str]								=							list(snake_case__						)

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

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

         lowerCAmelCase     :     List[Any]								=							[]
         for idx in range(0		,     len(snake_case__						)		,     snake_case__						):
           for l in iterable_dataset_lists:
             observed += l[idx : idx + shard_batch_size]

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

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     List[str]								=							42
         lowerCAmelCase     :     Tuple								=							RandomIterableDataset()

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

         # Edge case with a very small dataset
         lowerCAmelCase     :     List[str]								=							RandomIterableDataset(max_length=2						)

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

       def        lowercase__							(						self						):






         """simple docstring"""





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

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     List[str]								=							SkipDataLoader(list(range(16						)						)		,     batch_size=4		,     skip_batches=2						)
         self.assertListEqual([t.tolist() for t in dataloader]		,     [[8, 9, 10, 11], [12, 13, 14, 15]]						)

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     int								=							DataLoader(list(range(16						)						)		,     batch_size=4						)
         lowerCAmelCase     :     Optional[int]								=							skip_first_batches(snake_case__		,     num_batches=2						)
         self.assertListEqual([t.tolist() for t in new_dataloader]		,     [[8, 9, 10, 11], [12, 13, 14, 15]]						)

       def        lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase     :     Tuple								=							DataLoaderShard(list(range(16						)						)		,     batch_size=4						)
         for idx, _ in enumerate(snake_case__						):
           self.assertEqual(dataloader.end_of_dataloader		,     idx == 3						)

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







       def        lowercase__							(						self						):






         """simple docstring"""





         Accelerator()
         lowerCAmelCase     :     Dict								=							DataLoaderDispatcher(range(16						)		,     batch_size=4						)
         for idx, _ in enumerate(snake_case__						):
           self.assertEqual(dataloader.end_of_dataloader		,     idx == 3						)

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





