



import gc
import unittest

import numpy as np
import torch
import torch.nn.functional as F
from transformers import (
    ClapTextConfig,
    ClapTextModelWithProjection,
    RobertaTokenizer,
    SpeechTaHifiGan,
    SpeechTaHifiGanConfig,
)

from diffusers import (
    AudioLDMPipeline,
    AutoencoderKL,
    DDIMScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from ..pipeline_params import TEXT_TO_AUDIO_BATCH_PARAMS, TEXT_TO_AUDIO_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()






class 						lowercase		(					_A					,       unittest.TestCase):






	"""simple docstring"""


	a__ :   List[Any]						  =				AudioLDMPipeline
	a__ :   Union[str, Any]						  =				TEXT_TO_AUDIO_PARAMS
	a__ :   Any						  =				TEXT_TO_AUDIO_BATCH_PARAMS
	a__ :   str						  =				frozenset(
	    [
	        "num_inference_steps",
	        "num_waveforms_per_prompt",
	        "generator",
	        "latents",
	        "output_type",
	        "return_dict",
	        "callback",
	        "callback_steps",
	    ])




	def     _SCREAMING_SNAKE_CASE						(      self      :		List[Any]	)     ->					List[str]:
							torch.manual_seed(0	)
							UpperCAmelCase_=  UNetaDConditionModel(
							    block_out_channels=(32, 64)		,    layers_per_block=2		,    sample_size=32		,    in_channels=4		,    out_channels=4		,    down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")		,    up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")		,    cross_attention_dim=(32, 64)		,    class_embed_type="""simple_projection"""		,    projection_class_embeddings_input_dim=32		,    class_embeddings_concat=UpperCamelCase__		,    )
							UpperCAmelCase_=  DDIMScheduler(
							    beta_start=0.00_085		,    beta_end=0.012		,    beta_schedule="""scaled_linear"""		,    clip_sample=UpperCamelCase__		,    set_alpha_to_one=UpperCamelCase__		,    )
							torch.manual_seed(0	)
							UpperCAmelCase_=  AutoencoderKL(
							    block_out_channels=[32, 64]		,    in_channels=1		,    out_channels=1		,    down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]		,    up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]		,    latent_channels=4		,    )
							torch.manual_seed(0	)
							UpperCAmelCase_=  ClapTextConfig(
							    bos_token_id=0		,    eos_token_id=2		,    hidden_size=32		,    intermediate_size=37		,    layer_norm_eps=1E-05		,    num_attention_heads=4		,    num_hidden_layers=5		,    pad_token_id=1		,    vocab_size=1_000		,    projection_dim=32		,    )
							UpperCAmelCase_=  ClapTextModelWithProjection(UpperCamelCase__	)
							UpperCAmelCase_=  RobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-random-roberta"""		,    model_max_length=77	)

							UpperCAmelCase_=  SpeechTaHifiGanConfig(
							    model_in_dim=8		,    sampling_rate=16_000		,    upsample_initial_channel=16		,    upsample_rates=[2, 2]		,    upsample_kernel_sizes=[4, 4]		,    resblock_kernel_sizes=[3, 7]		,    resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5]]		,    normalize_before=UpperCamelCase__		,    )

							UpperCAmelCase_=  SpeechTaHifiGan(UpperCamelCase__	)

							UpperCAmelCase_=  {
							    """unet""": unet,
							    """scheduler""": scheduler,
							    """vae""": vae,
							    """text_encoder""": text_encoder,
							    """tokenizer""": tokenizer,
							    """vocoder""": vocoder,
							}
							return components




	def     _SCREAMING_SNAKE_CASE						(      self      :		Optional[int]		,    __UpperCAmelCase      :		str		,    __UpperCAmelCase      :		List[Any]=0	)     ->					Optional[Any]:






							if str(UpperCamelCase__	).startswith("""mps"""	):
													UpperCAmelCase_=  torch.manual_seed(UpperCamelCase__	)
							else:
													UpperCAmelCase_=  torch.Generator(device=UpperCamelCase__	).manual_seed(UpperCamelCase__	)
							UpperCAmelCase_=  {
							    """prompt""": """A hammer hitting a wooden surface""",
							    """generator""": generator,
							    """num_inference_steps""": 2,
							    """guidance_scale""": 6.0,
							}
							return inputs




	def     _SCREAMING_SNAKE_CASE						(      self      :		Tuple	)     ->					Optional[Any]:
							UpperCAmelCase_=  """cpu"""  # ensure determinism for the device-dependent torch.Generator

							UpperCAmelCase_=  self.get_dummy_components()
							UpperCAmelCase_=  AudioLDMPipeline(**UpperCamelCase__	)
							UpperCAmelCase_=  audioldm_pipe.to(UpperCamelCase__	)
							audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

							UpperCAmelCase_=  self.get_dummy_inputs(UpperCamelCase__	)
							UpperCAmelCase_=  audioldm_pipe(**UpperCamelCase__	)
							UpperCAmelCase_=  output.audios[0]

							assert audio.ndim == 1
							assert len(UpperCamelCase__	) == 256

							UpperCAmelCase_=  audio[:10]
							UpperCAmelCase_=  np.array(
							    [-0.0_050, 0.0_050, -0.0_060, 0.0_033, -0.0_026, 0.0_033, -0.0_027, 0.0_033, -0.0_028, 0.0_033]	)

							assert np.abs(audio_slice - expected_slice	).max() < 1E-2




	def     _SCREAMING_SNAKE_CASE						(      self      :		List[str]	)     ->					Any:
							UpperCAmelCase_=  self.get_dummy_components()
							UpperCAmelCase_=  AudioLDMPipeline(**UpperCamelCase__	)
							UpperCAmelCase_=  audioldm_pipe.to(UpperCamelCase__	)
							UpperCAmelCase_=  audioldm_pipe.to(UpperCamelCase__	)
							audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

							UpperCAmelCase_=  self.get_dummy_inputs(UpperCamelCase__	)
							UpperCAmelCase_=  3 * [inputs["""prompt"""]]

							# forward
							UpperCAmelCase_=  audioldm_pipe(**UpperCamelCase__	)
							UpperCAmelCase_=  output.audios[0]

							UpperCAmelCase_=  self.get_dummy_inputs(UpperCamelCase__	)
							UpperCAmelCase_=  3 * [inputs.pop("""prompt"""	)]

							UpperCAmelCase_=  audioldm_pipe.tokenizer(
							    UpperCamelCase__		,    padding="""max_length"""		,    max_length=audioldm_pipe.tokenizer.model_max_length		,    truncation=UpperCamelCase__		,    return_tensors="""pt"""		,    )
							UpperCAmelCase_=  text_inputs["""input_ids"""].to(UpperCamelCase__	)

							UpperCAmelCase_=  audioldm_pipe.text_encoder(
							    UpperCamelCase__		,    )
							UpperCAmelCase_=  prompt_embeds.text_embeds
							# additional L_2 normalization over each hidden-state
							UpperCAmelCase_=  F.normalize(UpperCamelCase__		,    dim=-1	)

							UpperCAmelCase_=  prompt_embeds

							# forward
							UpperCAmelCase_=  audioldm_pipe(**UpperCamelCase__	)
							UpperCAmelCase_=  output.audios[0]

							assert np.abs(audio_a - audio_a	).max() < 1E-2




	def     _SCREAMING_SNAKE_CASE						(      self      :		int	)     ->					str:
							UpperCAmelCase_=  self.get_dummy_components()
							UpperCAmelCase_=  AudioLDMPipeline(**UpperCamelCase__	)
							UpperCAmelCase_=  audioldm_pipe.to(UpperCamelCase__	)
							UpperCAmelCase_=  audioldm_pipe.to(UpperCamelCase__	)
							audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

							UpperCAmelCase_=  self.get_dummy_inputs(UpperCamelCase__	)
							UpperCAmelCase_=  3 * ["""this is a negative prompt"""]
							UpperCAmelCase_=  negative_prompt
							UpperCAmelCase_=  3 * [inputs["""prompt"""]]

							# forward
							UpperCAmelCase_=  audioldm_pipe(**UpperCamelCase__	)
							UpperCAmelCase_=  output.audios[0]

							UpperCAmelCase_=  self.get_dummy_inputs(UpperCamelCase__	)
							UpperCAmelCase_=  3 * [inputs.pop("""prompt"""	)]

							UpperCAmelCase_=  []
							for p in [prompt, negative_prompt]:
													UpperCAmelCase_=  audioldm_pipe.tokenizer(
													    UpperCamelCase__		,    padding="""max_length"""		,    max_length=audioldm_pipe.tokenizer.model_max_length		,    truncation=UpperCamelCase__		,    return_tensors="""pt"""		,    )
													UpperCAmelCase_=  text_inputs["""input_ids"""].to(UpperCamelCase__	)

													UpperCAmelCase_=  audioldm_pipe.text_encoder(
													    UpperCamelCase__		,    )
													UpperCAmelCase_=  text_embeds.text_embeds
													# additional L_2 normalization over each hidden-state
													UpperCAmelCase_=  F.normalize(UpperCamelCase__		,    dim=-1	)

													embeds.append(UpperCamelCase__	)

							UpperCAmelCase_,  UpperCAmelCase_=  embeds

							# forward
							UpperCAmelCase_=  audioldm_pipe(**UpperCamelCase__	)
							UpperCAmelCase_=  output.audios[0]

							assert np.abs(audio_a - audio_a	).max() < 1E-2




	def     _SCREAMING_SNAKE_CASE						(      self      :		Optional[Any]	)     ->					Optional[int]:
							UpperCAmelCase_=  """cpu"""  # ensure determinism for the device-dependent torch.Generator
							UpperCAmelCase_=  self.get_dummy_components()
							UpperCAmelCase_=  PNDMScheduler(skip_prk_steps=UpperCamelCase__	)
							UpperCAmelCase_=  AudioLDMPipeline(**UpperCamelCase__	)
							UpperCAmelCase_=  audioldm_pipe.to(UpperCamelCase__	)
							audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

							UpperCAmelCase_=  self.get_dummy_inputs(UpperCamelCase__	)
							UpperCAmelCase_=  """egg cracking"""
							UpperCAmelCase_=  audioldm_pipe(**UpperCamelCase__		,    negative_prompt=UpperCamelCase__	)
							UpperCAmelCase_=  output.audios[0]

							assert audio.ndim == 1
							assert len(UpperCamelCase__	) == 256

							UpperCAmelCase_=  audio[:10]
							UpperCAmelCase_=  np.array(
							    [-0.0_051, 0.0_050, -0.0_060, 0.0_034, -0.0_026, 0.0_033, -0.0_027, 0.0_033, -0.0_028, 0.0_032]	)

							assert np.abs(audio_slice - expected_slice	).max() < 1E-2




	def     _SCREAMING_SNAKE_CASE						(      self      :		Any	)     ->					Optional[Any]:
							UpperCAmelCase_=  """cpu"""  # ensure determinism for the device-dependent torch.Generator
							UpperCAmelCase_=  self.get_dummy_components()
							UpperCAmelCase_=  PNDMScheduler(skip_prk_steps=UpperCamelCase__	)
							UpperCAmelCase_=  AudioLDMPipeline(**UpperCamelCase__	)
							UpperCAmelCase_=  audioldm_pipe.to(UpperCamelCase__	)
							audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

							UpperCAmelCase_=  """A hammer hitting a wooden surface"""

							# test num_waveforms_per_prompt=1 (default)
							UpperCAmelCase_=  audioldm_pipe(UpperCamelCase__		,    num_inference_steps=2	).audios

							assert audios.shape == (1, 256)

							# test num_waveforms_per_prompt=1 (default) for batch of prompts
							UpperCAmelCase_=  2
							UpperCAmelCase_=  audioldm_pipe([prompt] * batch_size		,    num_inference_steps=2	).audios

							assert audios.shape == (batch_size, 256)

							# test num_waveforms_per_prompt for single prompt
							UpperCAmelCase_=  2
							UpperCAmelCase_=  audioldm_pipe(UpperCamelCase__		,    num_inference_steps=2		,    num_waveforms_per_prompt=UpperCamelCase__	).audios

							assert audios.shape == (num_waveforms_per_prompt, 256)

							# test num_waveforms_per_prompt for batch of prompts
							UpperCAmelCase_=  2
							UpperCAmelCase_=  audioldm_pipe(
							    [prompt] * batch_size		,    num_inference_steps=2		,    num_waveforms_per_prompt=UpperCamelCase__	).audios

							assert audios.shape == (batch_size * num_waveforms_per_prompt, 256)




	def     _SCREAMING_SNAKE_CASE						(      self      :		int	)     ->					Union[str, Any]:
							UpperCAmelCase_=  """cpu"""  # ensure determinism for the device-dependent torch.Generator
							UpperCAmelCase_=  self.get_dummy_components()
							UpperCAmelCase_=  AudioLDMPipeline(**UpperCamelCase__	)
							UpperCAmelCase_=  audioldm_pipe.to(UpperCamelCase__	)
							audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase__	)
							UpperCAmelCase_=  audioldm_pipe.vocoder.config.sampling_rate

							UpperCAmelCase_=  self.get_dummy_inputs(UpperCamelCase__	)
							UpperCAmelCase_=  audioldm_pipe(audio_length_in_s=0.016		,    **UpperCamelCase__	)
							UpperCAmelCase_=  output.audios[0]

							assert audio.ndim == 1
							assert len(UpperCamelCase__	) / vocoder_sampling_rate == 0.016

							UpperCAmelCase_=  audioldm_pipe(audio_length_in_s=0.032		,    **UpperCamelCase__	)
							UpperCAmelCase_=  output.audios[0]

							assert audio.ndim == 1
							assert len(UpperCamelCase__	) / vocoder_sampling_rate == 0.032




	def     _SCREAMING_SNAKE_CASE						(      self      :		Any	)     ->					Union[str, Any]:
							UpperCAmelCase_=  self.get_dummy_components()
							UpperCAmelCase_=  AudioLDMPipeline(**UpperCamelCase__	)
							UpperCAmelCase_=  audioldm_pipe.to(UpperCamelCase__	)
							audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

							UpperCAmelCase_=  ["""hey"""]

							UpperCAmelCase_=  audioldm_pipe(UpperCamelCase__		,    num_inference_steps=1	)
							UpperCAmelCase_=  output.audios.shape
							assert audio_shape == (1, 256)

							UpperCAmelCase_=  audioldm_pipe.vocoder.config
							config.model_in_dim *= 2
							UpperCAmelCase_=  SpeechTaHifiGan(UpperCamelCase__	).to(UpperCamelCase__	)
							UpperCAmelCase_=  audioldm_pipe(UpperCamelCase__		,    num_inference_steps=1	)
							UpperCAmelCase_=  output.audios.shape
							# waveform shape is unchanged, we just have 2x the number of mel channels in the spectrogram
							assert audio_shape == (1, 256)




	def     _SCREAMING_SNAKE_CASE						(      self      :		Tuple	)     ->					Dict:
							self._test_attention_slicing_forward_pass(test_mean_pixel_difference=UpperCamelCase__	)




	def     _SCREAMING_SNAKE_CASE						(      self      :		Optional[Any]	)     ->					List[str]:
							self._test_inference_batch_single_identical(test_mean_pixel_difference=UpperCamelCase__	)




	@unittest.skipIf(
	    torch_device != """cuda""" or not is_xformers_available()		,    reason="""XFormers attention is only available with CUDA and `xformers` installed"""		,    )
	def     _SCREAMING_SNAKE_CASE						(      self      :		Union[str, Any]	)     ->					Optional[Any]:
							self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=UpperCamelCase__	)






@slow
class 						lowercase		(					unittest.TestCase):






	"""simple docstring"""


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




	def     _SCREAMING_SNAKE_CASE						(      self      :		Dict		,    __UpperCAmelCase      :		Tuple		,    __UpperCAmelCase      :		Any="cpu"		,    __UpperCAmelCase      :		Optional[Any]=torch.floataa		,    __UpperCAmelCase      :		str=0	)     ->					Union[str, Any]:
							UpperCAmelCase_=  torch.Generator(device=UpperCamelCase__	).manual_seed(UpperCamelCase__	)
							UpperCAmelCase_=  np.random.RandomState(UpperCamelCase__	).standard_normal((1, 8, 128, 16)	)
							UpperCAmelCase_=  torch.from_numpy(UpperCamelCase__	).to(device=UpperCamelCase__		,    dtype=UpperCamelCase__	)
							UpperCAmelCase_=  {
							    """prompt""": """A hammer hitting a wooden surface""",
							    """latents""": latents,
							    """generator""": generator,
							    """num_inference_steps""": 3,
							    """guidance_scale""": 2.5,
							}
							return inputs




	def     _SCREAMING_SNAKE_CASE						(      self      :		Any	)     ->					List[str]:
							UpperCAmelCase_=  AudioLDMPipeline.from_pretrained("""cvssp/audioldm"""	)
							UpperCAmelCase_=  audioldm_pipe.to(UpperCamelCase__	)
							audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

							UpperCAmelCase_=  self.get_inputs(UpperCamelCase__	)
							UpperCAmelCase_=  25
							UpperCAmelCase_=  audioldm_pipe(**UpperCamelCase__	).audios[0]

							assert audio.ndim == 1
							assert len(UpperCamelCase__	) == 81_920

							UpperCAmelCase_=  audio[77_230:77_240]
							UpperCAmelCase_=  np.array(
							    [-0.4_884, -0.4_607, 0.0_023, 0.5_007, 0.5_896, 0.5_151, 0.3_813, -0.0_208, -0.3_687, -0.4_315]	)
							UpperCAmelCase_=  np.abs(expected_slice - audio_slice	).max()
							assert max_diff < 1E-2




	def     _SCREAMING_SNAKE_CASE						(      self      :		List[str]	)     ->					Dict:
							UpperCAmelCase_=  AudioLDMPipeline.from_pretrained("""cvssp/audioldm"""	)
							UpperCAmelCase_=  LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config	)
							UpperCAmelCase_=  audioldm_pipe.to(UpperCamelCase__	)
							audioldm_pipe.set_progress_bar_config(disable=UpperCamelCase__	)

							UpperCAmelCase_=  self.get_inputs(UpperCamelCase__	)
							UpperCAmelCase_=  audioldm_pipe(**UpperCamelCase__	).audios[0]

							assert audio.ndim == 1
							assert len(UpperCamelCase__	) == 81_920

							UpperCAmelCase_=  audio[27_780:27_790]
							UpperCAmelCase_=  np.array([-0.2_131, -0.0_873, -0.0_124, -0.0_189, 0.0_569, 0.1_373, 0.1_883, 0.2_886, 0.3_297, 0.2_212]	)
							UpperCAmelCase_=  np.abs(expected_slice - audio_slice	).max()
							assert max_diff < 3E-2






import random
import unittest

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

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




class 			UpperCAmelCase_     (					_A	):


  '''simple docstring'''







  def __init__(					self     :					Optional[int]     ,				UpperCamelCase__     :					Optional[Any]=0.01     ,				UpperCamelCase__     :					Tuple=1000       )     ->      str:







    """simple docstring"""






    __magic_name__								=			p_stop
    __magic_name__								=			max_length





  def __iter__(					self     :					Union[str, Any]       )     ->      Dict:







    """simple docstring"""






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







class 			UpperCAmelCase_     (					unittest.TestCase	):


  '''simple docstring'''







  def 				_lowercase				(					self     :					Optional[Any]     ,				UpperCamelCase__     :					Union[str, Any]     ,				UpperCamelCase__     :					Tuple     ,				UpperCamelCase__     :					str=False     ,				UpperCamelCase__     :					Optional[Any]=True       )     ->      Dict:







    """simple docstring"""






    __magic_name__								=			[
        BatchSamplerShard(UpperCamelCase__     ,				2     ,				UpperCamelCase__     ,				split_batches=UpperCamelCase__     ,				even_batches=UpperCamelCase__       )
        for i in range(2       )
    ]
    __magic_name__								=			[list(UpperCamelCase__       ) for batch_sampler_shard in batch_sampler_shards]
    if not split_batches:
      self.assertListEqual([len(UpperCamelCase__       ) for shard in batch_sampler_shards]     ,				[len(UpperCamelCase__       ) for e in expected]       )
    self.assertListEqual(UpperCamelCase__     ,				UpperCamelCase__       )
  def 				_lowercase				(					self     :					Any       )     ->      Union[str, Any]:







    """simple docstring"""






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

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

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

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

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

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

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

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

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

    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[], []]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__       )
  def 				_lowercase				(					self     :					Tuple       )     ->      int:







    """simple docstring"""






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

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

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

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

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

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

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

    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[], []]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				split_batches=UpperCamelCase__       )
  def 				_lowercase				(					self     :					int       )     ->      str:







    """simple docstring"""






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

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

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

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

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

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

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

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

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

    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=3     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[], []]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				even_batches=UpperCamelCase__       )
  def 				_lowercase				(					self     :					Optional[int]       )     ->      int:







    """simple docstring"""






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

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

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

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

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

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

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

    __magic_name__								=			BatchSampler(range(2       )     ,				batch_size=4     ,				drop_last=UpperCamelCase__       )
    __magic_name__								=			[[], []]
    self.check_batch_sampler_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				split_batches=UpperCamelCase__     ,				even_batches=UpperCamelCase__       )
  def 				_lowercase				(					self     :					Dict       )     ->      str:







    """simple docstring"""






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

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

    self.assertListEqual(list(batch_sampler_shards[0]       )     ,				[[0, 1, 2], [5, 6, 7, 8], [12, 13]]       )
    self.assertListEqual(list(batch_sampler_shards[1]       )     ,				[[3, 4], [9, 10, 11]]       )
  def 				_lowercase				(					self     :					int     ,				UpperCamelCase__     :					List[str]     ,				UpperCamelCase__     :					Union[str, Any]     ,				UpperCamelCase__     :					Tuple     ,				UpperCamelCase__     :					Dict=False     ,				UpperCamelCase__     :					Union[str, Any]=2     ,				UpperCamelCase__     :					Union[str, Any]=False       )     ->      Union[str, Any]:







    """simple docstring"""






    random.seed(UpperCamelCase__       )
    __magic_name__								=			list(UpperCamelCase__       )

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

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

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

    if not drop_last:
      while len(UpperCamelCase__       ) < len(UpperCamelCase__       ):
        reference += reference
    self.assertListEqual(UpperCamelCase__     ,				reference[: len(UpperCamelCase__       )]       )
  def 				_lowercase				(					self     :					Tuple       )     ->      List[str]:







    """simple docstring"""






    __magic_name__								=			42
    __magic_name__								=			RandomIterableDataset()

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

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

    self.check_iterable_dataset_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				batch_size=4     ,				drop_last=UpperCamelCase__     ,				split_batches=UpperCamelCase__       )
    self.check_iterable_dataset_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				batch_size=4     ,				drop_last=UpperCamelCase__     ,				split_batches=UpperCamelCase__       )
    self.check_iterable_dataset_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				batch_size=4     ,				drop_last=UpperCamelCase__     ,				split_batches=UpperCamelCase__       )
    self.check_iterable_dataset_shards(UpperCamelCase__     ,				UpperCamelCase__     ,				batch_size=4     ,				drop_last=UpperCamelCase__     ,				split_batches=UpperCamelCase__       )
  def 				_lowercase				(					self     :					List[Any]       )     ->      int:







    """simple docstring"""






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







    """simple docstring"""






    __magic_name__								=			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     :					str       )     ->      List[str]:







    """simple docstring"""






    __magic_name__								=			DataLoader(list(range(16       )       )     ,				batch_size=4       )
    __magic_name__								=			skip_first_batches(UpperCamelCase__     ,				num_batches=2       )
    self.assertListEqual([t.tolist() for t in new_dataloader]     ,				[[8, 9, 10, 11], [12, 13, 14, 15]]       )
  def 				_lowercase				(					self     :					List[str]       )     ->      Tuple:







    """simple docstring"""






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

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





  def 				_lowercase				(					self     :					Tuple       )     ->      List[str]:







    """simple docstring"""






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

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

