






'''simple docstring'''



from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


__UpperCamelCase   :							str         =   logging.get_logger(__name__)  # pylint: disable=invalid-name

__UpperCamelCase   :							str         =   '''
    Examples:
        ```py
        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior")
        >>> pipe_prior.to("cuda")
        >>> prompt = "red cat, 4k photo"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> zero_image_emb = out.negative_image_embeds
        >>> pipe = KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder")
        >>> pipe.to("cuda")
        >>> image = pipe(
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=50,
        ... ).images
        >>> image[0].save("cat.png")
        ```
'''



def  __UpperCAmelCase						(					SCREAMING_SNAKE_CASE__:						Optional[int],       SCREAMING_SNAKE_CASE__:						List[Any],       SCREAMING_SNAKE_CASE__:						Dict=8     )	->      List[str]:

   """simple docstring"""
   __a					  =     height // scale_factor**2
   if height % scale_factor**2 != 0:
      new_height += 1
   __a					  =     width // scale_factor**2
   if width % scale_factor**2 != 0:
      new_width += 1
   return new_height * scale_factor, new_width * scale_factor



class 				__SCREAMING_SNAKE_CASE							(				_lowerCAmelCase  ):




      def __init__(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, )				->Optional[int]:



         '''simple docstring'''
         super().__init__()

         self.register_modules(
             unet=_SCREAMING_SNAKE_CASE			, scheduler=_SCREAMING_SNAKE_CASE			, movq=_SCREAMING_SNAKE_CASE			, )
         __a					  =     2 ** (len(self.movq.config.block_out_channels   ) - 1)




      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )				->List[Any]:



         '''simple docstring'''





         if latents is None:
            __a					  =     randn_tensor(_SCREAMING_SNAKE_CASE			, generator=_SCREAMING_SNAKE_CASE			, device=_SCREAMING_SNAKE_CASE			, dtype=_SCREAMING_SNAKE_CASE   )
         else:
            if latents.shape != shape:
               raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {shape}"""   )
            __a					  =     latents.to(_SCREAMING_SNAKE_CASE   )

         __a					  =     latents * scheduler.init_noise_sigma
         return latents




      def 			__UpperCamelCase				(	self			, lowerCamelCase=0   )				->Union[str, Any]:



         '''simple docstring'''





         if is_accelerate_available():
            from accelerate import cpu_offload
         else:
            raise ImportError('Please install accelerate via `pip install accelerate`'   )

         __a					  =     torch.device(F"""cuda:{gpu_id}"""   )

         __a					  =     [
             self.unet,
             self.movq,
         ]
         for cpu_offloaded_model in models:
            if cpu_offloaded_model is not None:
               cpu_offload(_SCREAMING_SNAKE_CASE			, _SCREAMING_SNAKE_CASE   )




      def 			__UpperCamelCase				(	self			, lowerCamelCase=0   )				->Dict:



         '''simple docstring'''





         if is_accelerate_available() and is_accelerate_version('>='			, '0.17.0.dev0'   ):
            from accelerate import cpu_offload_with_hook
         else:
            raise ImportError('`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'   )

         __a					  =     torch.device(F"""cuda:{gpu_id}"""   )

         if self.device.type != "cpu":
            self.to('cpu'			, silence_dtype_warnings=_SCREAMING_SNAKE_CASE   )
            torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

         __a					  =     None
         for cpu_offloaded_model in [self.unet, self.movq]:
            __a  ,      __a					  =     cpu_offload_with_hook(_SCREAMING_SNAKE_CASE			, _SCREAMING_SNAKE_CASE			, prev_module_hook=_SCREAMING_SNAKE_CASE   )

         # We'll offload the last model manually.
         __a					  =     hook




      @property
      # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
      def 			__UpperCamelCase				(	self   )				->int:



         '''simple docstring'''





         if not hasattr(self.unet			, '_hf_hook'   ):
            return self.device
         for module in self.unet.modules():
            if (
                hasattr(_SCREAMING_SNAKE_CASE			, '_hf_hook'   )
                and hasattr(module._hf_hook			, 'execution_device'   )
                and module._hf_hook.execution_device is not None
            ):
               return torch.device(module._hf_hook.execution_device   )
         return self.device






      @torch.no_grad()
      @replace_example_docstring(_SCREAMING_SNAKE_CASE   )
      def __call__(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase = 512			, lowerCamelCase = 512			, lowerCamelCase = 100			, lowerCamelCase = 4.0			, lowerCamelCase = 1			, lowerCamelCase = None			, lowerCamelCase = None			, lowerCamelCase = "pil"			, lowerCamelCase = True			, )				->Optional[Any]:



         '''simple docstring'''
         __a					  =     self._execution_device

         __a					  =     guidance_scale > 1.0

         if isinstance(_SCREAMING_SNAKE_CASE			, _SCREAMING_SNAKE_CASE   ):
            __a					  =     torch.cat(_SCREAMING_SNAKE_CASE			, dim=0   )
         __a					  =     image_embeds.shape[0] * num_images_per_prompt
         if isinstance(_SCREAMING_SNAKE_CASE			, _SCREAMING_SNAKE_CASE   ):
            __a					  =     torch.cat(_SCREAMING_SNAKE_CASE			, dim=0   )

         if do_classifier_free_guidance:
            __a					  =     image_embeds.repeat_interleave(_SCREAMING_SNAKE_CASE			, dim=0   )
            __a					  =     negative_image_embeds.repeat_interleave(_SCREAMING_SNAKE_CASE			, dim=0   )

         __a					  =     torch.cat([negative_image_embeds, image_embeds]			, dim=0   ).to(dtype=self.unet.dtype			, device=_SCREAMING_SNAKE_CASE   )

         self.scheduler.set_timesteps(_SCREAMING_SNAKE_CASE			, device=_SCREAMING_SNAKE_CASE   )
         __a					  =     self.scheduler.timesteps

         __a					  =     self.unet.config.in_channels

         __a  ,      __a					  =     downscale_height_and_width(_SCREAMING_SNAKE_CASE			, _SCREAMING_SNAKE_CASE			, self.movq_scale_factor   )

         # create initial latent
         __a					  =     self.prepare_latents(
             (batch_size, num_channels_latents, height, width)			, image_embeds.dtype			, _SCREAMING_SNAKE_CASE			, _SCREAMING_SNAKE_CASE			, _SCREAMING_SNAKE_CASE			, self.scheduler			, )

         for i, t in enumerate(self.progress_bar(_SCREAMING_SNAKE_CASE   )   ):
            # expand the latents if we are doing classifier free guidance
            __a					  =     torch.cat([latents] * 2   ) if do_classifier_free_guidance else latents

            __a					  =     {'image_embeds': image_embeds}
            __a					  =     self.unet(
                sample=_SCREAMING_SNAKE_CASE			, timestep=_SCREAMING_SNAKE_CASE			, encoder_hidden_states=_SCREAMING_SNAKE_CASE			, added_cond_kwargs=_SCREAMING_SNAKE_CASE			, return_dict=_SCREAMING_SNAKE_CASE			, )[0]

            if do_classifier_free_guidance:
               __a  ,      __a					  =     noise_pred.split(latents.shape[1]			, dim=1   )
               __a  ,      __a					  =     noise_pred.chunk(2   )
               __a  ,      __a					  =     variance_pred.chunk(2   )
               __a					  =     noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
               __a					  =     torch.cat([noise_pred, variance_pred_text]			, dim=1   )

            if not (
                hasattr(self.scheduler.config			, 'variance_type'   )
                and self.scheduler.config.variance_type in ["learned", "learned_range"]
            ):
               __a  ,      __a					  =     noise_pred.split(latents.shape[1]			, dim=1   )

            # compute the previous noisy sample x_t -> x_t-1
            __a					  =     self.scheduler.step(
                _SCREAMING_SNAKE_CASE			, _SCREAMING_SNAKE_CASE			, _SCREAMING_SNAKE_CASE			, generator=_SCREAMING_SNAKE_CASE			, )[0]
         # post-processing
         __a					  =     self.movq.decode(_SCREAMING_SNAKE_CASE			, force_not_quantize=_SCREAMING_SNAKE_CASE   )['sample']

         if output_type not in ["pt", "np", "pil"]:
            raise ValueError(F"""Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}"""   )

         if output_type in ["np", "pil"]:
            __a					  =     image * 0.5 + 0.5
            __a					  =     image.clamp(0			, 1   )
            __a					  =     image.cpu().permute(0			, 2			, 3			, 1   ).float().numpy()

         if output_type == "pil":
            __a					  =     self.numpy_to_pil(_SCREAMING_SNAKE_CASE   )

         if not return_dict:
            return (image,)

         return ImagePipelineOutput(images=_SCREAMING_SNAKE_CASE   )



import random
import unittest

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

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

class 						__a (							UpperCAmelCase     ):
		def __init__(     self     ,       _SCREAMING_SNAKE_CASE=0.01     ,       _SCREAMING_SNAKE_CASE=1000		)					->			str:




									"""simple docstring"""





									_UpperCAmelCase											=     p_stop
									_UpperCAmelCase											=     max_length






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




									"""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 						__a (							unittest.TestCase     ):
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE=False     ,       _SCREAMING_SNAKE_CASE=True		)					->			int:




									"""simple docstring"""





									_UpperCAmelCase											=     [
									    BatchSamplerShard(_SCREAMING_SNAKE_CASE     ,       2     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)
									    for i in range(2		)
									]
									_UpperCAmelCase											=     [list(_SCREAMING_SNAKE_CASE		) for batch_sampler_shard in batch_sampler_shards]
									if not split_batches:
																self.assertListEqual([len(_SCREAMING_SNAKE_CASE		) for shard in batch_sampler_shards]     ,       [len(_SCREAMING_SNAKE_CASE		) for e in expected]		)
									self.assertListEqual(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     BatchSampler(range(24		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_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(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

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

									# 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=_SCREAMING_SNAKE_CASE		)
									_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(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

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

									# 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=_SCREAMING_SNAKE_CASE		)
									_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(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(22		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_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.
									_UpperCAmelCase											=     BatchSampler(range(20		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_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(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)

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

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

									_UpperCAmelCase											=     BatchSampler(range(2		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [[], []]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Union[str, Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     BatchSampler(range(24		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_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(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)

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

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

									_UpperCAmelCase											=     BatchSampler(range(22		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_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(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)

									# 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=_SCREAMING_SNAKE_CASE		)
									_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(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(21		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_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(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)

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

									_UpperCAmelCase											=     BatchSampler(range(2		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [[], []]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Tuple:




									"""simple docstring"""





									_UpperCAmelCase											=     BatchSampler(range(24		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_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(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

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

									# 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=_SCREAMING_SNAKE_CASE		)
									_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(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

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

									# 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=_SCREAMING_SNAKE_CASE		)
									_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(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(22		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [
									    [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
									    [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
									]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_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.
									_UpperCAmelCase											=     BatchSampler(range(20		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_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(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

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

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

									_UpperCAmelCase											=     BatchSampler(range(2		)     ,       batch_size=3     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [[], []]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     BatchSampler(range(24		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_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(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

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

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

									_UpperCAmelCase											=     BatchSampler(range(22		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_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(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									# 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=_SCREAMING_SNAKE_CASE		)
									_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(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     BatchSampler(range(21		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_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(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)

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

									_UpperCAmelCase											=     BatchSampler(range(2		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     [[], []]
									self.check_batch_sampler_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
									_UpperCAmelCase											=     [BatchSamplerShard(_SCREAMING_SNAKE_CASE     ,       2     ,       _SCREAMING_SNAKE_CASE     ,       even_batches=_SCREAMING_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       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE=False     ,       _SCREAMING_SNAKE_CASE=2     ,       _SCREAMING_SNAKE_CASE=False		)					->			Dict:




									"""simple docstring"""





									random.seed(_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     list(_SCREAMING_SNAKE_CASE		)

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

									_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(_SCREAMING_SNAKE_CASE		)     ,       len(_SCREAMING_SNAKE_CASE		)		)
																self.assertTrue(len(_SCREAMING_SNAKE_CASE		) % shard_batch_size == 0		)

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

									if not drop_last:
																while len(_SCREAMING_SNAKE_CASE		) < len(_SCREAMING_SNAKE_CASE		):
																							reference += reference
									self.assertListEqual(_SCREAMING_SNAKE_CASE     ,       reference[: len(_SCREAMING_SNAKE_CASE		)]		)
		def       UpperCAmelCase__							(     self		)					->			List[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     42
									_UpperCAmelCase											=     RandomIterableDataset()

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

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

									self.check_iterable_dataset_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)
									self.check_iterable_dataset_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)
									self.check_iterable_dataset_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)
									self.check_iterable_dataset_shards(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE     ,       split_batches=_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			str:




									"""simple docstring"""





									_UpperCAmelCase											=     BatchSampler(range(16		)     ,       batch_size=4     ,       drop_last=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     SkipBatchSampler(_SCREAMING_SNAKE_CASE     ,       2		)
									self.assertListEqual(list(_SCREAMING_SNAKE_CASE		)     ,       [[8, 9, 10, 11], [12, 13, 14, 15]]		)
		def       UpperCAmelCase__							(     self		)					->			Tuple:




									"""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       UpperCAmelCase__							(     self		)					->			int:




									"""simple docstring"""





									_UpperCAmelCase											=     DataLoader(list(range(16		)		)     ,       batch_size=4		)
									_UpperCAmelCase											=     skip_first_batches(_SCREAMING_SNAKE_CASE     ,       num_batches=2		)
									self.assertListEqual([t.tolist() for t in new_dataloader]     ,       [[8, 9, 10, 11], [12, 13, 14, 15]]		)
		def       UpperCAmelCase__							(     self		)					->			int:




									"""simple docstring"""





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

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






		def       UpperCAmelCase__							(     self		)					->			Any:




									"""simple docstring"""





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

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