




'''simple docstring'''


import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference





class  __lowercase    (	__lowercase							,				unittest.TestCase			):
   _a    						=   ShapEImgaImgPipeline
   _a    						=   ["image"]
   _a    						=   ["image"]
   _a    						=   [
       "num_images_per_prompt",
       "num_inference_steps",
       "generator",
       "latents",
       "guidance_scale",
       "frame_size",
       "output_type",
       "return_dict",
   ]
   _a    						=   False







   @property
   def 		UpperCamelCase__      (						self   )					->     Dict:
     return 32







   @property
   def 		UpperCamelCase__      (						self   )					->     Optional[int]:
     return 32







   @property
   def 		UpperCamelCase__      (						self   )					->     Optional[int]:
     return self.time_input_dim * 4







   @property
   def 		UpperCamelCase__      (						self   )					->     str:
     return 8







   @property
   def 		UpperCamelCase__      (						self   )					->     Any:
     torch.manual_seed(0   )
     __a								=       CLIPVisionConfig(
         hidden_size=self.text_embedder_hidden_size   ,				image_size=64   ,				projection_dim=self.text_embedder_hidden_size   ,				intermediate_size=37   ,				num_attention_heads=4   ,				num_channels=3   ,				num_hidden_layers=5   ,				patch_size=1   ,				)

     __a								=       CLIPVisionModel(UpperCamelCase   )
     return model







   @property
   def 		UpperCamelCase__      (						self   )					->     Optional[Any]:
     __a								=       CLIPImageProcessor(
         crop_size=224   ,				do_center_crop=UpperCamelCase   ,				do_normalize=UpperCamelCase   ,				do_resize=UpperCamelCase   ,				image_mean=[0.48_145_466, 0.4_578_275, 0.40_821_073]   ,				image_std=[0.26_862_954, 0.26_130_258, 0.27_577_711]   ,				resample=3   ,				size=224   ,				)

     return image_processor







   @property
   def 		UpperCamelCase__      (						self   )					->     List[Any]:
     torch.manual_seed(0   )

     __a								=       {
         'num_attention_heads': 2,
         'attention_head_dim': 16,
         'embedding_dim': self.time_input_dim,
         'num_embeddings': 32,
         'embedding_proj_dim': self.text_embedder_hidden_size,
         'time_embed_dim': self.time_embed_dim,
         'num_layers': 1,
         'clip_embed_dim': self.time_input_dim * 2,
         'additional_embeddings': 0,
         'time_embed_act_fn': 'gelu',
         'norm_in_type': 'layer',
         'embedding_proj_norm_type': 'layer',
         'encoder_hid_proj_type': None,
         'added_emb_type': None,
     }

     __a								=       PriorTransformer(**UpperCamelCase   )
     return model







   @property
   def 		UpperCamelCase__      (						self   )					->     Optional[int]:
     torch.manual_seed(0   )

     __a								=       {
         'param_shapes': (
             (self.renderer_dim, 93),
             (self.renderer_dim, 8),
             (self.renderer_dim, 8),
             (self.renderer_dim, 8),
         ),
         'd_latent': self.time_input_dim,
         'd_hidden': self.renderer_dim,
         'n_output': 12,
         'background': (
             0.1,
             0.1,
             0.1,
         ),
     }
     __a								=       ShapERenderer(**UpperCamelCase   )
     return model







   def 		UpperCamelCase__      (						self   )					->     Any:
     __a								=       self.dummy_prior
     __a								=       self.dummy_image_encoder
     __a								=       self.dummy_image_processor
     __a								=       self.dummy_renderer

     __a								=       HeunDiscreteScheduler(
         beta_schedule='exp'   ,				num_train_timesteps=1024   ,				prediction_type='sample'   ,				use_karras_sigmas=UpperCamelCase   ,				clip_sample=UpperCamelCase   ,				clip_sample_range=1.0   ,				)
     __a								=       {
         'prior': prior,
         'image_encoder': image_encoder,
         'image_processor': image_processor,
         'renderer': renderer,
         'scheduler': scheduler,
     }

     return components







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase=0   )					->     Optional[int]:
     __a								=       floats_tensor((1, 3, 64, 64)   ,				rng=random.Random(UpperCamelCase   )   ).to(UpperCamelCase   )

     if str(UpperCamelCase   ).startswith('mps'   ):
       __a								=       torch.manual_seed(UpperCamelCase   )
     else:
       __a								=       torch.Generator(device=UpperCamelCase   ).manual_seed(UpperCamelCase   )
     __a								=       {
         'image': input_image,
         'generator': generator,
         'num_inference_steps': 1,
         'frame_size': 32,
         'output_type': 'np',
     }
     return inputs







   def 		UpperCamelCase__      (						self   )					->     Dict:
     __a								=       'cpu'

     __a								=       self.get_dummy_components()

     __a								=       self.pipeline_class(**UpperCamelCase   )
     __a								=       pipe.to(UpperCamelCase   )

     pipe.set_progress_bar_config(disable=UpperCamelCase   )

     __a								=       pipe(**self.get_dummy_inputs(UpperCamelCase   )   )
     __a								=       output.images[0]
     __a								=       image[0, -3:, -3:, -1]

     assert image.shape == (20, 32, 32, 3)

     __a								=       np.array(
         [
             0.00_039_216,
             0.00_039_216,
             0.00_039_216,
             0.00_039_216,
             0.00_039_216,
             0.00_039_216,
             0.00_039_216,
             0.00_039_216,
             0.00_039_216,
         ]   )

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







   def 		UpperCamelCase__      (						self   )					->     Optional[Any]:
     self._test_inference_batch_consistent(batch_sizes=[1, 2]   )







   def 		UpperCamelCase__      (						self   )					->     Tuple:
     __a								=       torch_device == 'cpu'
     __a								=       True
     self._test_inference_batch_single_identical(
         batch_size=2   ,				test_max_difference=UpperCamelCase   ,				relax_max_difference=UpperCamelCase   ,				)







   def 		UpperCamelCase__      (						self   )					->     Tuple:
     __a								=       self.get_dummy_components()
     __a								=       self.pipeline_class(**UpperCamelCase   )
     __a								=       pipe.to(UpperCamelCase   )
     pipe.set_progress_bar_config(disable=UpperCamelCase   )

     __a								=       1
     __a								=       2

     __a								=       self.get_dummy_inputs(UpperCamelCase   )

     for key in inputs.keys():
       if key in self.batch_params:
         __a								=       batch_size * [inputs[key]]

     __a								=       pipe(**UpperCamelCase   ,				num_images_per_prompt=UpperCamelCase   )[0]

     assert images.shape[0] == batch_size * num_images_per_prompt





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







   def 		UpperCamelCase__      (						self   )					->     List[Any]:
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()







   def 		UpperCamelCase__      (						self   )					->     str:
     __a								=       load_image(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/shap_e/corgi.png'   )
     __a								=       load_numpy(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
         '/shap_e/test_shap_e_img2img_out.npy'   )
     __a								=       ShapEImgaImgPipeline.from_pretrained('openai/shap-e-img2img'   )
     __a								=       pipe.to(UpperCamelCase   )
     pipe.set_progress_bar_config(disable=UpperCamelCase   )

     __a								=       torch.Generator(device=UpperCamelCase   ).manual_seed(0   )

     __a								=       pipe(
         UpperCamelCase   ,				generator=UpperCamelCase   ,				guidance_scale=3.0   ,				num_inference_steps=64   ,				frame_size=64   ,				output_type='np'   ,				).images[0]

     assert images.shape == (20, 64, 64, 3)

     assert_mean_pixel_difference(UpperCamelCase   ,				UpperCamelCase   )





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 							lowerCAmelCase__   ( __lowercase						):

						def __init__(      self							,				a=0.01							,				a=10_00					)       ->   Optional[int]:


								'''simple docstring'''

								_UpperCamelCase         =						p_stop
								_UpperCamelCase         =						max_length







						def __iter__(      self					)       ->   Optional[Any]:


								'''simple docstring'''

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







class 							lowerCAmelCase__   ( unittest.TestCase						):

						def 			A_					(      self							,				a							,				a							,				a=False							,				a=True					)       ->   Optional[Any]:


								'''simple docstring'''

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

						def 			A_					(      self					)       ->   Any:


								'''simple docstring'''

								_UpperCamelCase         =						BatchSampler(range(24					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
								    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
								]
								self.check_batch_sampler_shards(a							,				a					)

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

								# Check the shards when the dataset is a round multiple of batch size but not total batch size.
								_UpperCamelCase         =						BatchSampler(range(21					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
								    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [0, 1, 2]],
								]
								self.check_batch_sampler_shards(a							,				a					)

								_UpperCamelCase         =						BatchSampler(range(21					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
								    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
								]
								self.check_batch_sampler_shards(a							,				a					)

								# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
								# num_processes batch.
								_UpperCamelCase         =						BatchSampler(range(22					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
								    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 0, 1]],
								]
								self.check_batch_sampler_shards(a							,				a					)

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

								# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
								# num_processes batch.
								_UpperCamelCase         =						BatchSampler(range(20					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 0]],
								    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [1, 2, 3]],
								]
								self.check_batch_sampler_shards(a							,				a					)

								_UpperCamelCase         =						BatchSampler(range(20					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
								    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
								]
								self.check_batch_sampler_shards(a							,				a					)

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

								_UpperCamelCase         =						BatchSampler(range(2					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[[], []]
								self.check_batch_sampler_shards(a							,				a					)

						def 			A_					(      self					)       ->   Tuple:


								'''simple docstring'''

								_UpperCamelCase         =						BatchSampler(range(24					)							,				batch_size=4							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
								    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
								]
								self.check_batch_sampler_shards(a							,				a							,				split_batches=a					)

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

								# Check the shards when the dataset is not a round multiple of batch size.
								_UpperCamelCase         =						BatchSampler(range(22					)							,				batch_size=4							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
								    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [0, 1]],
								]
								self.check_batch_sampler_shards(a							,				a							,				split_batches=a					)

								_UpperCamelCase         =						BatchSampler(range(22					)							,				batch_size=4							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
								    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
								]
								self.check_batch_sampler_shards(a							,				a							,				split_batches=a					)

								# Check the shards when the dataset is not a round multiple of batch size or num_processes.
								_UpperCamelCase         =						BatchSampler(range(21					)							,				batch_size=4							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 0]],
								    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [1, 2]],
								]
								self.check_batch_sampler_shards(a							,				a							,				split_batches=a					)

								_UpperCamelCase         =						BatchSampler(range(21					)							,				batch_size=4							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
								    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
								]
								self.check_batch_sampler_shards(a							,				a							,				split_batches=a					)

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

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

						def 			A_					(      self					)       ->   str:


								'''simple docstring'''

								_UpperCamelCase         =						BatchSampler(range(24					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
								    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
								]
								self.check_batch_sampler_shards(a							,				a							,				even_batches=a					)

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

								# Check the shards when the dataset is a round multiple of batch size but not total batch size.
								_UpperCamelCase         =						BatchSampler(range(21					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
								    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
								]
								self.check_batch_sampler_shards(a							,				a							,				even_batches=a					)

								_UpperCamelCase         =						BatchSampler(range(21					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
								    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
								]
								self.check_batch_sampler_shards(a							,				a							,				even_batches=a					)

								# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
								# num_processes batch.
								_UpperCamelCase         =						BatchSampler(range(22					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
								    [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21]],
								]
								self.check_batch_sampler_shards(a							,				a							,				even_batches=a					)

								_UpperCamelCase         =						BatchSampler(range(22					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
								    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
								]
								self.check_batch_sampler_shards(a							,				a							,				even_batches=a					)

								# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
								# num_processes batch.
								_UpperCamelCase         =						BatchSampler(range(20					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19]],
								    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
								]
								self.check_batch_sampler_shards(a							,				a							,				even_batches=a					)

								_UpperCamelCase         =						BatchSampler(range(20					)							,				batch_size=3							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
								    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
								]
								self.check_batch_sampler_shards(a							,				a							,				even_batches=a					)

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

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

						def 			A_					(      self					)       ->   Optional[Any]:


								'''simple docstring'''

								_UpperCamelCase         =						BatchSampler(range(24					)							,				batch_size=4							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
								    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
								]
								self.check_batch_sampler_shards(a							,				a							,				split_batches=a							,				even_batches=a					)

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

								# Check the shards when the dataset is not a round multiple of batch size.
								_UpperCamelCase         =						BatchSampler(range(22					)							,				batch_size=4							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
								    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
								]
								self.check_batch_sampler_shards(a							,				a							,				split_batches=a							,				even_batches=a					)

								_UpperCamelCase         =						BatchSampler(range(22					)							,				batch_size=4							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
								    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
								]
								self.check_batch_sampler_shards(a							,				a							,				split_batches=a							,				even_batches=a					)

								# Check the shards when the dataset is not a round multiple of batch size or num_processes.
								_UpperCamelCase         =						BatchSampler(range(21					)							,				batch_size=4							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20]],
								    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
								]
								self.check_batch_sampler_shards(a							,				a							,				split_batches=a							,				even_batches=a					)

								_UpperCamelCase         =						BatchSampler(range(21					)							,				batch_size=4							,				drop_last=a					)
								_UpperCamelCase         =						[
								    [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
								    [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
								]
								self.check_batch_sampler_shards(a							,				a							,				split_batches=a							,				even_batches=a					)

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

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

						def 			A_					(      self					)       ->   List[Any]:


								'''simple docstring'''

								_UpperCamelCase         =						[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
								_UpperCamelCase         =						[BatchSamplerShard(a							,				2							,				a							,				even_batches=a					) for i in range(2					)]

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

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

						def 			A_					(      self							,				a							,				a							,				a							,				a=False							,				a=2							,				a=False					)       ->   Any:


								'''simple docstring'''

								random.seed(a					)
								_UpperCamelCase         =						list(a					)

								_UpperCamelCase         =						[
								    IterableDatasetShard(
								        a							,				batch_size=a							,				drop_last=a							,				num_processes=a							,				process_index=a							,				split_batches=a							,				)
								    for i in range(a					)
								]
								_UpperCamelCase         =						[]
								for iterable_dataset_shard in iterable_dataset_shards:
										# Since our random iterable dataset will be... random... we need to use a seed to get reproducible results.
										random.seed(a					)
										iterable_dataset_lists.append(list(a					)					)

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

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

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

						def 			A_					(      self					)       ->   int:


								'''simple docstring'''

								_UpperCamelCase         =						42
								_UpperCamelCase         =						RandomIterableDataset()

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

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

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

						def 			A_					(      self					)       ->   Any:


								'''simple docstring'''

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

						def 			A_					(      self					)       ->   Dict:


								'''simple docstring'''

								_UpperCamelCase         =						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 			A_					(      self					)       ->   Tuple:


								'''simple docstring'''

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

						def 			A_					(      self					)       ->   List[str]:


								'''simple docstring'''

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

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







						def 			A_					(      self					)       ->   Optional[Any]:


								'''simple docstring'''

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

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



