



import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionInstructPixaPixPipeline,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils import floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class 						__a (     SCREAMING_SNAKE_CASE	,    SCREAMING_SNAKE_CASE	,    SCREAMING_SNAKE_CASE	,    unittest.TestCase				):
	SCREAMING_SNAKE_CASE       =       StableDiffusionInstructPixaPixPipeline
	SCREAMING_SNAKE_CASE       =       TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width", "cross_attention_kwargs"}
	SCREAMING_SNAKE_CASE       =       TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
	SCREAMING_SNAKE_CASE       =       IMAGE_TO_IMAGE_IMAGE_PARAMS
	SCREAMING_SNAKE_CASE       =       IMAGE_TO_IMAGE_IMAGE_PARAMS



	def 		UpperCamelCase							(						self							:    str)->							List[str]:
						torch.manual_seed(0)
						__lowerCAmelCase	   =UNetaDConditionModel(
						    block_out_channels=(32, 64)				,	layers_per_block=2				,	sample_size=32				,	in_channels=8				,	out_channels=4				,	down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")				,	up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")				,	cross_attention_dim=32				,	)
						__lowerCAmelCase	   =PNDMScheduler(skip_prk_steps=snake_case_)
						torch.manual_seed(0)
						__lowerCAmelCase	   =AutoencoderKL(
						    block_out_channels=[32, 64]				,	in_channels=3				,	out_channels=3				,	down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]				,	up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]				,	latent_channels=4				,	)
						torch.manual_seed(0)
						__lowerCAmelCase	   =CLIPTextConfig(
						    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=10_00				,	)
						__lowerCAmelCase	   =CLIPTextModel(snake_case_)
						__lowerCAmelCase	   =CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""")

						__lowerCAmelCase	   ={
						    """unet""": unet,
						    """scheduler""": scheduler,
						    """vae""": vae,
						    """text_encoder""": text_encoder,
						    """tokenizer""": tokenizer,
						    """safety_checker""": None,
						    """feature_extractor""": None,
						}
						return components



	def 		UpperCamelCase							(						self							:    Any				,	snake_case_							:    Union[str, Any]				,	snake_case_							:    str=0)->							int:
						__lowerCAmelCase	   =floats_tensor((1, 3, 32, 32)				,	rng=random.Random(snake_case_)).to(snake_case_)
						__lowerCAmelCase	   =image.cpu().permute(0				,	2				,	3				,	1)[0]
						__lowerCAmelCase	   =Image.fromarray(np.uinta(snake_case_)).convert("""RGB""")
						if str(snake_case_).startswith("""mps"""):
											__lowerCAmelCase	   =torch.manual_seed(snake_case_)
						else:
											__lowerCAmelCase	   =torch.Generator(device=snake_case_).manual_seed(snake_case_)
						__lowerCAmelCase	   ={
						    """prompt""": """A painting of a squirrel eating a burger""",
						    """image""": image,
						    """generator""": generator,
						    """num_inference_steps""": 2,
						    """guidance_scale""": 6.0,
						    """image_guidance_scale""": 1,
						    """output_type""": """numpy""",
						}
						return inputs



	def 		UpperCamelCase							(						self							:    Dict)->							str:
						__lowerCAmelCase	   ="""cpu"""  # ensure determinism for the device-dependent torch.Generator
						__lowerCAmelCase	   =self.get_dummy_components()
						__lowerCAmelCase	   =StableDiffusionInstructPixaPixPipeline(**snake_case_)
						__lowerCAmelCase	   =sd_pipe.to(snake_case_)
						sd_pipe.set_progress_bar_config(disable=snake_case_)

						__lowerCAmelCase	   =self.get_dummy_inputs(snake_case_)
						__lowerCAmelCase	   =sd_pipe(**snake_case_).images
						__lowerCAmelCase	   =image[0, -3:, -3:, -1]
						assert image.shape == (1, 32, 32, 3)
						__lowerCAmelCase	   =np.array([0.7_5_2_6, 0.3_7_5_0, 0.4_5_4_7, 0.6_1_1_7, 0.5_8_6_6, 0.5_0_1_6, 0.4_3_2_7, 0.5_6_4_2, 0.4_8_1_5])

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



	def 		UpperCamelCase							(						self							:    Optional[Any])->							str:
						__lowerCAmelCase	   ="""cpu"""  # ensure determinism for the device-dependent torch.Generator
						__lowerCAmelCase	   =self.get_dummy_components()
						__lowerCAmelCase	   =StableDiffusionInstructPixaPixPipeline(**snake_case_)
						__lowerCAmelCase	   =sd_pipe.to(snake_case_)
						sd_pipe.set_progress_bar_config(disable=snake_case_)

						__lowerCAmelCase	   =self.get_dummy_inputs(snake_case_)
						__lowerCAmelCase	   ="""french fries"""
						__lowerCAmelCase	   =sd_pipe(**snake_case_				,	negative_prompt=snake_case_)
						__lowerCAmelCase	   =output.images
						__lowerCAmelCase	   =image[0, -3:, -3:, -1]

						assert image.shape == (1, 32, 32, 3)
						__lowerCAmelCase	   =np.array([0.7_5_1_1, 0.3_6_4_2, 0.4_5_5_3, 0.6_2_3_6, 0.5_7_9_7, 0.5_0_1_3, 0.4_3_4_3, 0.5_6_1_1, 0.4_8_3_1])

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



	def 		UpperCamelCase							(						self							:    Any)->							Optional[int]:
						__lowerCAmelCase	   ="""cpu"""  # ensure determinism for the device-dependent torch.Generator
						__lowerCAmelCase	   =self.get_dummy_components()
						__lowerCAmelCase	   =StableDiffusionInstructPixaPixPipeline(**snake_case_)
						__lowerCAmelCase	   =sd_pipe.to(snake_case_)
						sd_pipe.set_progress_bar_config(disable=snake_case_)

						__lowerCAmelCase	   =self.get_dummy_inputs(snake_case_)
						__lowerCAmelCase	   =[inputs["""prompt"""]] * 2

						__lowerCAmelCase	   =np.array(inputs["""image"""]).astype(np.floataa) / 2_5_5.0
						__lowerCAmelCase	   =torch.from_numpy(snake_case_).unsqueeze(0).to(snake_case_)
						__lowerCAmelCase	   =image / 2 + 0.5
						__lowerCAmelCase	   =image.permute(0				,	3				,	1				,	2)
						__lowerCAmelCase	   =image.repeat(2				,	1				,	1				,	1)

						__lowerCAmelCase	   =sd_pipe(**snake_case_).images
						__lowerCAmelCase	   =image[-1, -3:, -3:, -1]

						assert image.shape == (2, 32, 32, 3)
						__lowerCAmelCase	   =np.array([0.5_8_1_2, 0.5_7_4_8, 0.5_2_2_2, 0.5_9_0_8, 0.5_6_9_5, 0.7_1_7_4, 0.6_8_0_4, 0.5_5_2_3, 0.5_5_7_9])

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



	def 		UpperCamelCase							(						self							:    Tuple)->							Any:
						__lowerCAmelCase	   ="""cpu"""  # ensure determinism for the device-dependent torch.Generator
						__lowerCAmelCase	   =self.get_dummy_components()
						__lowerCAmelCase	   =EulerAncestralDiscreteScheduler(
						    beta_start=0.0_0_0_8_5				,	beta_end=0.0_1_2				,	beta_schedule="""scaled_linear""")
						__lowerCAmelCase	   =StableDiffusionInstructPixaPixPipeline(**snake_case_)
						__lowerCAmelCase	   =sd_pipe.to(snake_case_)
						sd_pipe.set_progress_bar_config(disable=snake_case_)

						__lowerCAmelCase	   =self.get_dummy_inputs(snake_case_)
						__lowerCAmelCase	   =sd_pipe(**snake_case_).images
						__lowerCAmelCase	   =image[0, -3:, -3:, -1]

						__lowerCAmelCase	   =[round(snake_case_				,	4) for x in image_slice.flatten().tolist()]
						print(""",""".join([str(snake_case_) for x in slice]))

						assert image.shape == (1, 32, 32, 3)
						__lowerCAmelCase	   =np.array([0.7_4_1_7, 0.3_8_4_2, 0.4_7_3_2, 0.5_7_7_6, 0.5_8_9_1, 0.5_1_3_9, 0.4_0_5_2, 0.5_6_7_3, 0.4_9_8_6])

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



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



	def 		UpperCamelCase							(						self							:    List[Any])->							Optional[Any]:
						__lowerCAmelCase	   =self.get_dummy_components()
						__lowerCAmelCase	   =StableDiffusionInstructPixaPixPipeline(**snake_case_)
						__lowerCAmelCase	   =VaeImageProcessor(do_resize=snake_case_				,	do_normalize=snake_case_)
						__lowerCAmelCase	   =pipe.to(snake_case_)
						pipe.set_progress_bar_config(disable=snake_case_)

						__lowerCAmelCase	   =pipe(**self.get_dummy_inputs_by_type(snake_case_				,	input_image_type="""pt"""))[0]

						__lowerCAmelCase	   =components["""vae"""]
						__lowerCAmelCase	   =self.get_dummy_inputs_by_type(snake_case_				,	input_image_type="""pt""")

						for image_param in self.image_latents_params:
											if image_param in inputs.keys():
																__lowerCAmelCase	   =vae.encode(inputs[image_param]).latent_dist.mode()

						__lowerCAmelCase	   =pipe(**snake_case_)[0]

						__lowerCAmelCase	   =np.abs(out - out_latents_inputs).max()
						self.assertLess(snake_case_				,	1e-4				,	"""passing latents as image input generate different result from passing image""")






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



	def 		UpperCamelCase							(						self							:    str)->							Union[str, Any]:
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()



	def 		UpperCamelCase							(						self							:    str				,	snake_case_							:    Any=0)->							List[Any]:
						__lowerCAmelCase	   =torch.manual_seed(snake_case_)
						__lowerCAmelCase	   =load_image(
						    """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg""")
						__lowerCAmelCase	   ={
						    """prompt""": """turn him into a cyborg""",
						    """image""": image,
						    """generator""": generator,
						    """num_inference_steps""": 3,
						    """guidance_scale""": 7.5,
						    """image_guidance_scale""": 1.0,
						    """output_type""": """numpy""",
						}
						return inputs



	def 		UpperCamelCase							(						self							:    Optional[Any])->							Union[str, Any]:
						__lowerCAmelCase	   =StableDiffusionInstructPixaPixPipeline.from_pretrained(
						    """timbrooks/instruct-pix2pix"""				,	safety_checker=snake_case_)
						pipe.to(snake_case_)
						pipe.set_progress_bar_config(disable=snake_case_)
						pipe.enable_attention_slicing()

						__lowerCAmelCase	   =self.get_inputs()
						__lowerCAmelCase	   =pipe(**snake_case_).images
						__lowerCAmelCase	   =image[0, -3:, -3:, -1].flatten()

						assert image.shape == (1, 5_12, 5_12, 3)
						__lowerCAmelCase	   =np.array([0.5_9_0_2, 0.6_0_1_5, 0.6_0_2_7, 0.5_9_8_3, 0.6_0_9_2, 0.6_0_6_1, 0.5_7_6_5, 0.5_7_8_5, 0.5_5_5_5])

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



	def 		UpperCamelCase							(						self							:    Any)->							Union[str, Any]:
						__lowerCAmelCase	   =StableDiffusionInstructPixaPixPipeline.from_pretrained(
						    """timbrooks/instruct-pix2pix"""				,	safety_checker=snake_case_)
						__lowerCAmelCase	   =LMSDiscreteScheduler.from_config(pipe.scheduler.config)
						pipe.to(snake_case_)
						pipe.set_progress_bar_config(disable=snake_case_)
						pipe.enable_attention_slicing()

						__lowerCAmelCase	   =self.get_inputs()
						__lowerCAmelCase	   =pipe(**snake_case_).images
						__lowerCAmelCase	   =image[0, -3:, -3:, -1].flatten()

						assert image.shape == (1, 5_12, 5_12, 3)
						__lowerCAmelCase	   =np.array([0.6_5_7_8, 0.6_8_1_7, 0.6_9_7_2, 0.6_7_6_1, 0.6_8_5_6, 0.6_9_1_6, 0.6_4_2_8, 0.6_5_1_6, 0.6_3_0_1])

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



	def 		UpperCamelCase							(						self							:    Optional[int])->							Union[str, Any]:
						__lowerCAmelCase	   =StableDiffusionInstructPixaPixPipeline.from_pretrained(
						    """timbrooks/instruct-pix2pix"""				,	safety_checker=snake_case_)
						__lowerCAmelCase	   =DDIMScheduler.from_config(pipe.scheduler.config)
						pipe.to(snake_case_)
						pipe.set_progress_bar_config(disable=snake_case_)
						pipe.enable_attention_slicing()

						__lowerCAmelCase	   =self.get_inputs()
						__lowerCAmelCase	   =pipe(**snake_case_).images
						__lowerCAmelCase	   =image[0, -3:, -3:, -1].flatten()

						assert image.shape == (1, 5_12, 5_12, 3)
						__lowerCAmelCase	   =np.array([0.3_8_2_8, 0.3_8_3_4, 0.3_8_1_8, 0.3_7_9_2, 0.3_8_6_5, 0.3_7_5_2, 0.3_7_9_2, 0.3_8_4_7, 0.3_7_5_3])

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



	def 		UpperCamelCase							(						self							:    List[Any])->							Tuple:
						__lowerCAmelCase	   =0

						def callback_fn(snake_case_							:    int				,	snake_case_							:    int				,	snake_case_							:    torch.FloatTensor) -> None:
											__lowerCAmelCase	   =True
											nonlocal number_of_steps
											number_of_steps += 1
											if step == 1:
																__lowerCAmelCase	   =latents.detach().cpu().numpy()
																assert latents.shape == (1, 4, 64, 64)
																__lowerCAmelCase	   =latents[0, -3:, -3:, -1]
																__lowerCAmelCase	   =np.array([-0.2_4_6_3, -0.4_6_4_4, -0.9_7_5_6, 1.5_1_7_6, 1.4_4_1_4, 0.7_8_6_6, 0.9_8_9_7, 0.8_5_2_1, 0.7_9_8_3])

																assert np.abs(latents_slice.flatten() - expected_slice).max() < 5e-2
											elif step == 2:
																__lowerCAmelCase	   =latents.detach().cpu().numpy()
																assert latents.shape == (1, 4, 64, 64)
																__lowerCAmelCase	   =latents[0, -3:, -3:, -1]
																__lowerCAmelCase	   =np.array([-0.2_6_4_4, -0.4_6_2_6, -0.9_6_5_3, 1.5_1_7_6, 1.4_5_5_1, 0.7_6_8_6, 0.9_8_0_5, 0.8_4_5_2, 0.8_1_1_5])

																assert np.abs(latents_slice.flatten() - expected_slice).max() < 5e-2

						__lowerCAmelCase	   =False

						__lowerCAmelCase	   =StableDiffusionInstructPixaPixPipeline.from_pretrained(
						    """timbrooks/instruct-pix2pix"""				,	safety_checker=snake_case_				,	torch_dtype=torch.floataa)
						__lowerCAmelCase	   =pipe.to(snake_case_)
						pipe.set_progress_bar_config(disable=snake_case_)
						pipe.enable_attention_slicing()

						__lowerCAmelCase	   =self.get_inputs()
						pipe(**snake_case_				,	callback=snake_case_				,	callback_steps=1)
						assert callback_fn.has_been_called
						assert number_of_steps == 3



	def 		UpperCamelCase							(						self							:    Optional[Any])->							Dict:
						torch.cuda.empty_cache()
						torch.cuda.reset_max_memory_allocated()
						torch.cuda.reset_peak_memory_stats()

						__lowerCAmelCase	   =StableDiffusionInstructPixaPixPipeline.from_pretrained(
						    """timbrooks/instruct-pix2pix"""				,	safety_checker=snake_case_				,	torch_dtype=torch.floataa)
						__lowerCAmelCase	   =pipe.to(snake_case_)
						pipe.set_progress_bar_config(disable=snake_case_)
						pipe.enable_attention_slicing(1)
						pipe.enable_sequential_cpu_offload()

						__lowerCAmelCase	   =self.get_inputs()
						__lowerCAmelCase	   =pipe(**snake_case_)

						__lowerCAmelCase	   =torch.cuda.max_memory_allocated()
						# make sure that less than 2.2 GB is allocated
						assert mem_bytes < 2.2 * 10**9



	def 		UpperCamelCase							(						self							:    List[Any])->							List[str]:
						__lowerCAmelCase	   =self.get_inputs()
						# resize to resolution that is divisible by 8 but not 16 or 32
						__lowerCAmelCase	   =inputs["""image"""].resize((5_04, 5_04))

						__lowerCAmelCase	   ="""timbrooks/instruct-pix2pix"""
						__lowerCAmelCase	   =StableDiffusionInstructPixaPixPipeline.from_pretrained(
						    snake_case_				,	safety_checker=snake_case_				,	)
						pipe.to(snake_case_)
						pipe.set_progress_bar_config(disable=snake_case_)
						pipe.enable_attention_slicing()

						__lowerCAmelCase	   =pipe(**snake_case_)
						__lowerCAmelCase	   =output.images[0]

						__lowerCAmelCase	   =image[2_55:2_58, 3_83:3_86, -1]

						assert image.shape == (5_04, 5_04, 3)
						__lowerCAmelCase	   =np.array([0.2_7_2_6, 0.2_5_2_9, 0.2_6_6_4, 0.2_6_5_5, 0.2_6_4_1, 0.2_6_4_2, 0.2_5_9_1, 0.2_6_4_9, 0.2_5_9_0])

						assert np.abs(image_slice.flatten() - expected_slice).max() < 5e-3









import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch






def 					__lowerCAmelCase		(  __lowerCamelCase	:    List[Any]	)  ->						int:  # picklable for multiprocessing
					return x.sum()






def 					__lowerCAmelCase		(  __lowerCamelCase	:    Any	)  ->						Any:  # picklable for multiprocessing
					return i + 1






@dataclass
class 						__a :
	SCREAMING_SNAKE_CASE       =       42
	SCREAMING_SNAKE_CASE       =       42






class 						__a (     SCREAMING_SNAKE_CASE				):



	def 		UpperCamelCase							(						self							:    List[Any])->							Tuple:
						__lowerCAmelCase	   ={}
						__lowerCAmelCase	   =[]
						__lowerCAmelCase	   =1
						__lowerCAmelCase	   =[1, 2]
						__lowerCAmelCase	   ={"""a""": 1, """b""": 2}
						__lowerCAmelCase	   ={"""a""": [1, 2], """b""": [3, 4]}
						__lowerCAmelCase	   ={"""a""": {"""1""": 1}, """b""": 2}
						__lowerCAmelCase	   ={"""a""": 1, """b""": 2, """c""": 3, """d""": 4}
						__lowerCAmelCase	   ={}
						__lowerCAmelCase	   =[]
						__lowerCAmelCase	   =2
						__lowerCAmelCase	   =[2, 3]
						__lowerCAmelCase	   ={"""a""": 2, """b""": 3}
						__lowerCAmelCase	   ={"""a""": [2, 3], """b""": [4, 5]}
						__lowerCAmelCase	   ={"""a""": {"""1""": 2}, """b""": 3}
						__lowerCAmelCase	   ={"""a""": 2, """b""": 3, """c""": 4, """d""": 5}
						self.assertEqual(map_nested(snake_case_				,	snake_case_)				,	snake_case_)
						self.assertEqual(map_nested(snake_case_				,	snake_case_)				,	snake_case_)
						self.assertEqual(map_nested(snake_case_				,	snake_case_)				,	snake_case_)
						self.assertEqual(map_nested(snake_case_				,	snake_case_)				,	snake_case_)
						self.assertEqual(map_nested(snake_case_				,	snake_case_)				,	snake_case_)
						self.assertEqual(map_nested(snake_case_				,	snake_case_)				,	snake_case_)
						self.assertEqual(map_nested(snake_case_				,	snake_case_)				,	snake_case_)
						self.assertEqual(map_nested(snake_case_				,	snake_case_)				,	snake_case_)

						__lowerCAmelCase	   =2
						self.assertEqual(map_nested(snake_case_				,	snake_case_				,	num_proc=snake_case_)				,	snake_case_)
						self.assertEqual(map_nested(snake_case_				,	snake_case_				,	num_proc=snake_case_)				,	snake_case_)
						self.assertEqual(map_nested(snake_case_				,	snake_case_				,	num_proc=snake_case_)				,	snake_case_)
						self.assertEqual(map_nested(snake_case_				,	snake_case_				,	num_proc=snake_case_)				,	snake_case_)
						self.assertEqual(map_nested(snake_case_				,	snake_case_				,	num_proc=snake_case_)				,	snake_case_)
						self.assertEqual(map_nested(snake_case_				,	snake_case_				,	num_proc=snake_case_)				,	snake_case_)
						self.assertEqual(map_nested(snake_case_				,	snake_case_				,	num_proc=snake_case_)				,	snake_case_)
						self.assertEqual(map_nested(snake_case_				,	snake_case_				,	num_proc=snake_case_)				,	snake_case_)

						__lowerCAmelCase	   ={"""a""": np.eye(2), """b""": np.zeros(3), """c""": np.ones(2)}
						__lowerCAmelCase	   ={"""a""": 2, """b""": 0, """c""": 2}
						__lowerCAmelCase	   ={
						    """a""": np.eye(2).astype(snake_case_),
						    """b""": np.zeros(3).astype(snake_case_),
						    """c""": np.ones(2).astype(snake_case_),
						}
						self.assertEqual(map_nested(snake_case_				,	snake_case_				,	map_numpy=snake_case_)				,	snake_case_)
						self.assertEqual(
						    {k: v.tolist() for k, v in map_nested(snake_case_				,	snake_case_				,	map_numpy=snake_case_).items()}				,	{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}				,	)
						self.assertEqual(map_nested(snake_case_				,	snake_case_				,	map_numpy=snake_case_				,	num_proc=snake_case_)				,	snake_case_)
						self.assertEqual(
						    {k: v.tolist() for k, v in map_nested(snake_case_				,	snake_case_				,	map_numpy=snake_case_				,	num_proc=snake_case_).items()}				,	{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}				,	)
						with self.assertRaises(snake_case_):  # can't pickle a local lambda
											map_nested(lambda snake_case_: x + 1				,	snake_case_				,	num_proc=snake_case_)



	def 		UpperCamelCase							(						self							:    Any)->							int:
						__lowerCAmelCase	   ={"""a""": 1, """b""": 2}
						__lowerCAmelCase	   ={"""a""": 3, """b""": 4}
						__lowerCAmelCase	   ={"""a""": 5, """b""": 6}
						__lowerCAmelCase	   =sorted([("""a""", (1, 3, 5)), ("""b""", (2, 4, 6))])
						self.assertEqual(sorted(zip_dict(snake_case_				,	snake_case_				,	snake_case_))				,	snake_case_)



	def 		UpperCamelCase							(						self							:    Optional[Any])->							Optional[int]:






						class 						__a :
							SCREAMING_SNAKE_CASE       =       "bar"


						__lowerCAmelCase	   =Foo()
						self.assertEqual(foo.my_attr				,	"""bar""")
						with temporary_assignment(snake_case_				,	"""my_attr"""				,	"""BAR"""):
											self.assertEqual(foo.my_attr				,	"""BAR""")
						self.assertEqual(foo.my_attr				,	"""bar""")






@pytest.mark.parametrize(
    """iterable_length, num_proc, expected_num_proc"""				,   [
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ]				,   )
def 					__lowerCAmelCase		(  __lowerCamelCase	:    List[str]				,   __lowerCamelCase	:    Any				,   __lowerCamelCase	:    Dict	)  ->						Union[str, Any]:
					with patch("""datasets.utils.py_utils._single_map_nested"""	) as mock_single_map_nested, patch(
					    """datasets.parallel.parallel.Pool"""	) as mock_multiprocessing_pool:
										__lowerCAmelCase	   ={f"""{i}""": i for i in range(__lowerCamelCase	)}
										__lowerCAmelCase	   =map_nested(lambda __lowerCamelCase	: x + 10				,   __lowerCamelCase				,   num_proc=__lowerCamelCase				,   parallel_min_length=16	)
										if expected_num_proc == 1:
															assert mock_single_map_nested.called
															assert not mock_multiprocessing_pool.called
										else:
															assert not mock_single_map_nested.called
															assert mock_multiprocessing_pool.called
															assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc






class 						__a (     SCREAMING_SNAKE_CASE				):



	@require_tf
	def 		UpperCamelCase							(						self							:    Tuple)->							Optional[int]:
						import tensorflow as tf
						from tensorflow.keras import layers

						__lowerCAmelCase	   =layers.Dense(2)

						def gen_random_output():
											__lowerCAmelCase	   =tf.random.uniform((1, 3))
											return model(snake_case_).numpy()

						with temp_seed(42				,	set_tensorflow=snake_case_):
											__lowerCAmelCase	   =gen_random_output()
						with temp_seed(42				,	set_tensorflow=snake_case_):
											__lowerCAmelCase	   =gen_random_output()
						__lowerCAmelCase	   =gen_random_output()

						np.testing.assert_equal(snake_case_				,	snake_case_)
						self.assertGreater(np.abs(outa - outa).sum()				,	0)



	@require_torch
	def 		UpperCamelCase							(						self							:    Dict)->							Dict:
						import torch

						def gen_random_output():
											__lowerCAmelCase	   =torch.nn.Linear(3				,	2)
											__lowerCAmelCase	   =torch.rand(1				,	3)
											return model(snake_case_).detach().numpy()

						with temp_seed(42				,	set_pytorch=snake_case_):
											__lowerCAmelCase	   =gen_random_output()
						with temp_seed(42				,	set_pytorch=snake_case_):
											__lowerCAmelCase	   =gen_random_output()
						__lowerCAmelCase	   =gen_random_output()

						np.testing.assert_equal(snake_case_				,	snake_case_)
						self.assertGreater(np.abs(outa - outa).sum()				,	0)



	def 		UpperCamelCase							(						self							:    Union[str, Any])->							List[Any]:
						def gen_random_output():
											return np.random.rand(1				,	3)

						with temp_seed(42):
											__lowerCAmelCase	   =gen_random_output()
						with temp_seed(42):
											__lowerCAmelCase	   =gen_random_output()
						__lowerCAmelCase	   =gen_random_output()

						np.testing.assert_equal(snake_case_				,	snake_case_)
						self.assertGreater(np.abs(outa - outa).sum()				,	0)






@pytest.mark.parametrize("""input_data"""				,   [{}]	)
def 					__lowerCAmelCase		(  __lowerCamelCase	:    List[str]	)  ->						List[str]:
					__lowerCAmelCase	   =NestedDataStructure(__lowerCamelCase	).data
					assert output_data == input_data






@pytest.mark.parametrize(
    """data, expected_output"""				,   [
        ({}, []),
        ([], []),
        ("""foo""", ["""foo"""]),
        (["""foo""", """bar"""], ["""foo""", """bar"""]),
        ([["""foo""", """bar"""]], ["""foo""", """bar"""]),
        ([[["""foo"""], ["""bar"""]]], ["""foo""", """bar"""]),
        ([[["""foo"""], """bar"""]], ["""foo""", """bar"""]),
        ({"""a""": 1, """b""": 2}, [1, 2]),
        ({"""a""": [1, 2], """b""": [3, 4]}, [1, 2, 3, 4]),
        ({"""a""": [[1, 2]], """b""": [[3, 4]]}, [1, 2, 3, 4]),
        ({"""a""": [[1, 2]], """b""": [3, 4]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [[[3], [4]]]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [[3, 4]]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [3, 4]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [3, [4]]}, [1, 2, 3, 4]),
        ({"""a""": {"""1""": 1}, """b""": 2}, [1, 2]),
        ({"""a""": {"""1""": [1]}, """b""": 2}, [1, 2]),
        ({"""a""": {"""1""": [1]}, """b""": [2]}, [1, 2]),
    ]				,   )
def 					__lowerCAmelCase		(  __lowerCamelCase	:    List[Any]				,   __lowerCamelCase	:    int	)  ->						List[str]:
					__lowerCAmelCase	   =NestedDataStructure(__lowerCamelCase	).flatten()
					assert output == expected_output






def 					__lowerCAmelCase		(  )  ->						int:
					__lowerCAmelCase	   =A(x=1				,   y="""foobar"""	)
					__lowerCAmelCase	   ={"""x""": 1, """y""": """foobar"""}
					assert asdict(__lowerCamelCase	) == expected_output

					__lowerCAmelCase	   ={"""a""": {"""b""": A(x=10				,   y="""foo"""	)}, """c""": [A(x=20				,   y="""bar"""	)]}
					__lowerCAmelCase	   ={"""a""": {"""b""": {"""x""": 10, """y""": """foo"""}}, """c""": [{"""x""": 20, """y""": """bar"""}]}
					assert asdict(__lowerCamelCase	) == expected_output

					with pytest.raises(__lowerCamelCase	):
										asdict([1, A(x=10				,   y="""foo"""	)]	)






def 					__lowerCAmelCase		(  __lowerCamelCase	:    str	)  ->						Dict:
					return text.split()






def 					__lowerCAmelCase		(  __lowerCamelCase	:    List[Any]	)  ->						Any:
					yield (time.time(), content)
					time.sleep(2	)
					yield (time.time(), content)






def 					__lowerCAmelCase		(  )  ->						Any:
					with Pool(2	) as pool:
										__lowerCAmelCase	   =list(iflatmap_unordered(__lowerCamelCase				,   _split_text				,   kwargs_iterable=[{"""text""": """hello there"""}] * 10	)	)
										assert out.count("""hello"""	) == 10
										assert out.count("""there"""	) == 10
										assert len(__lowerCamelCase	) == 20

					# check multiprocess from pathos (uses dill for pickling)
					with multiprocess.Pool(2	) as pool:
										__lowerCAmelCase	   =list(iflatmap_unordered(__lowerCamelCase				,   _split_text				,   kwargs_iterable=[{"""text""": """hello there"""}] * 10	)	)
										assert out.count("""hello"""	) == 10
										assert out.count("""there"""	) == 10
										assert len(__lowerCamelCase	) == 20

					# check that we get items as fast as possible
					with Pool(2	) as pool:
										__lowerCAmelCase	   =[]
										for yield_time, content in iflatmap_unordered(
										    __lowerCamelCase				,   _aseconds_generator_of_aitems_with_timing				,   kwargs_iterable=[{"""content""": """a"""}, {"""content""": """b"""}]	):
															assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
															out.append(__lowerCamelCase	)
										assert out.count("""a"""	) == 2
										assert out.count("""b"""	) == 2
										assert len(__lowerCamelCase	) == 4




