





from typing import List, Optional, Union

import torch
from transformers import (
    XLMRobertaTokenizer,
)

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


__lowerCamelCase :  Optional[Any]										=  logging.get_logger(__name__)  # pylint: disable=invalid-name

__lowerCamelCase :  List[Any]										=  """
    Examples:
        ```py
        >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyPriorPipeline.from_pretrained(\"kandinsky-community/Kandinsky-2-1-prior\")
        >>> pipe_prior.to(\"cuda\")

        >>> prompt = \"red cat, 4k photo\"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> negative_image_emb = out.negative_image_embeds

        >>> pipe = KandinskyPipeline.from_pretrained(\"kandinsky-community/kandinsky-2-1\")
        >>> pipe.to(\"cuda\")

        >>> image = pipe(
        ...     prompt,
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=negative_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=100,
        ... ).images

        >>> image[0].save(\"cat.png\")
        ```
"""



def 	A_       (  _lowerCAmelCase							,     _lowerCAmelCase							,     _lowerCAmelCase=8  ):
 UpperCamelCase			:      Union[str, Any] 			=     h // scale_factor**2
 if h % scale_factor**2 != 0:
  new_h += 1
 UpperCamelCase			:      Any 			=     w // scale_factor**2
 if w % scale_factor**2 != 0:
  new_w += 1
 return new_h * scale_factor, new_w * scale_factor



class       A__     ( __snake_case    ):





      def __init__(				self ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							):
       '''simple docstring'''




       super().__init__()

       self.register_modules(
           text_encoder=A_ ,							tokenizer=A_ ,							unet=A_ ,							scheduler=A_ ,							movq=A_ ,							)
       UpperCamelCase			:      List[Any] 			=     2 ** (len(self.movq.config.block_out_channels     ) - 1)





      def  __UpperCamelCase(				self ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       if latents is None:
        UpperCamelCase			:      int 			=     randn_tensor(A_ ,							generator=A_ ,							device=A_ ,							dtype=A_     )
       else:
        if latents.shape != shape:
         raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {shape}"""     )
        UpperCamelCase			:      Optional[Any] 			=     latents.to(A_     )

       UpperCamelCase			:      Dict 			=     latents * scheduler.init_noise_sigma
       return latents





      def  __UpperCamelCase(				self ,							A_ ,							A_ ,							A_ ,							A_ ,							A_=None ,							):
       '''simple docstring'''




       UpperCamelCase			:      int 			=     len(A_     ) if isinstance(A_ ,							A_     ) else 1
       # get prompt text embeddings
       UpperCamelCase			:      Any 			=     self.tokenizer(
           A_ ,							padding="max_length" ,							truncation=A_ ,							max_length=77 ,							return_attention_mask=A_ ,							add_special_tokens=A_ ,							return_tensors="pt" ,							)

       UpperCamelCase			:      List[Any] 			=     text_inputs.input_ids
       UpperCamelCase			:      List[Any] 			=     self.tokenizer(A_ ,							padding="longest" ,							return_tensors="pt"     ).input_ids

       if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(A_ ,							A_     ):
        UpperCamelCase			:      Union[str, Any] 			=     self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]     )
        logger.warning(
            "The following part of your input was truncated because CLIP can only handle sequences up to"
            F""" {self.tokenizer.model_max_length} tokens: {removed_text}"""     )

       UpperCamelCase			:      str 			=     text_input_ids.to(A_     )
       UpperCamelCase			:      int 			=     text_inputs.attention_mask.to(A_     )

       UpperCamelCase			:      Optional[int] 			=     self.text_encoder(
           input_ids=A_ ,							attention_mask=A_     )

       UpperCamelCase			:      Optional[Any] 			=     prompt_embeds.repeat_interleave(A_ ,							dim=0     )
       UpperCamelCase			:      Any 			=     text_encoder_hidden_states.repeat_interleave(A_ ,							dim=0     )
       UpperCamelCase			:      int 			=     text_mask.repeat_interleave(A_ ,							dim=0     )

       if do_classifier_free_guidance:
        UpperCamelCase			:      List[str]
        if negative_prompt is None:
         UpperCamelCase			:      Optional[Any] 			=     [""] * batch_size
        elif type(A_     ) is not type(A_     ):
         raise TypeError(
             F"""`negative_prompt` should be the same type to `prompt`, but got {type(A_     )} !="""
             F""" {type(A_     )}."""     )
        elif isinstance(A_ ,							A_     ):
         UpperCamelCase			:      Union[str, Any] 			=     [negative_prompt]
        elif batch_size != len(A_     ):
         raise ValueError(
             F"""`negative_prompt`: {negative_prompt} has batch size {len(A_     )}, but `prompt`:"""
             F""" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"""
             " the batch size of `prompt`."     )
        else:
         UpperCamelCase			:      Optional[int] 			=     negative_prompt

        UpperCamelCase			:      Tuple 			=     self.tokenizer(
            A_ ,							padding="max_length" ,							max_length=77 ,							truncation=A_ ,							return_attention_mask=A_ ,							add_special_tokens=A_ ,							return_tensors="pt" ,							)
        UpperCamelCase			:      Optional[Any] 			=     uncond_input.input_ids.to(A_     )
        UpperCamelCase			:      Tuple 			=     uncond_input.attention_mask.to(A_     )

        UpperCamelCase			:      Optional[int] 			=     self.text_encoder(
            input_ids=A_ ,							attention_mask=A_     )

        # duplicate unconditional embeddings for each generation per prompt, using mps friendly method

        UpperCamelCase			:      int 			=     negative_prompt_embeds.shape[1]
        UpperCamelCase			:      Any 			=     negative_prompt_embeds.repeat(1 ,							A_     )
        UpperCamelCase			:      List[Any] 			=     negative_prompt_embeds.view(batch_size * num_images_per_prompt ,							A_     )

        UpperCamelCase			:      List[Any] 			=     uncond_text_encoder_hidden_states.shape[1]
        UpperCamelCase			:      int 			=     uncond_text_encoder_hidden_states.repeat(1 ,							A_ ,							1     )
        UpperCamelCase			:      Union[str, Any] 			=     uncond_text_encoder_hidden_states.view(
            batch_size * num_images_per_prompt ,							A_ ,							-1     )
        UpperCamelCase			:      Optional[int] 			=     uncond_text_mask.repeat_interleave(A_ ,							dim=0     )

        # done duplicates

        # For classifier free guidance, we need to do two forward passes.
        # Here we concatenate the unconditional and text embeddings into a single batch
        # to avoid doing two forward passes
        UpperCamelCase			:      List[Any] 			=     torch.cat([negative_prompt_embeds, prompt_embeds]     )
        UpperCamelCase			:      Optional[int] 			=     torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]     )

        UpperCamelCase			:      Tuple 			=     torch.cat([uncond_text_mask, text_mask]     )

       return prompt_embeds, text_encoder_hidden_states, text_mask





      def  __UpperCamelCase(				self ,							A_=0     ):
       '''simple docstring'''




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

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

       UpperCamelCase			:      Optional[int] 			=     [
           self.unet,
           self.text_encoder,
           self.movq,
       ]
       for cpu_offloaded_model in models:
        if cpu_offloaded_model is not None:
         cpu_offload(A_ ,							A_     )





      def  __UpperCamelCase(				self ,							A_=0     ):
       '''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."     )

       UpperCamelCase			:      Union[str, Any] 			=     torch.device(F"""cuda:{gpu_id}"""     )

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

       UpperCamelCase			:      Tuple 			=     None
       for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
        UpperCamelCase			:      Optional[int] 			=     cpu_offload_with_hook(A_ ,							A_ ,							prev_module_hook=A_     )

       if self.safety_checker is not None:
        UpperCamelCase			:      Any 			=     cpu_offload_with_hook(self.safety_checker ,							A_ ,							prev_module_hook=A_     )

       # We'll offload the last model manually.
       UpperCamelCase			:      str 			=     hook





      @property
      # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       if not hasattr(self.unet ,							"_hf_hook"     ):
        return self.device
       for module in self.unet.modules():
        if (
            hasattr(A_ ,							"_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(A_     )
      def __call__(				self ,							A_ ,							A_ ,							A_ ,							A_ = None ,							A_ = 512 ,							A_ = 512 ,							A_ = 100 ,							A_ = 4.0 ,							A_ = 1 ,							A_ = None ,							A_ = None ,							A_ = "pil" ,							A_ = True ,							):
       '''simple docstring'''




       if isinstance(A_ ,							A_     ):
        UpperCamelCase			:      Optional[Any] 			=     1
       elif isinstance(A_ ,							A_     ):
        UpperCamelCase			:      List[Any] 			=     len(A_     )
       else:
        raise ValueError(F"""`prompt` has to be of type `str` or `list` but is {type(A_     )}"""     )

       UpperCamelCase			:      str 			=     self._execution_device

       UpperCamelCase			:      Dict 			=     batch_size * num_images_per_prompt
       UpperCamelCase			:      Optional[int] 			=     guidance_scale > 1.0

       UpperCamelCase			:      Union[str, Any] 			=     self._encode_prompt(
           A_ ,							A_ ,							A_ ,							A_ ,							A_     )

       if isinstance(A_ ,							A_     ):
        UpperCamelCase			:      Any 			=     torch.cat(A_ ,							dim=0     )
       if isinstance(A_ ,							A_     ):
        UpperCamelCase			:      int 			=     torch.cat(A_ ,							dim=0     )

       if do_classifier_free_guidance:
        UpperCamelCase			:      int 			=     image_embeds.repeat_interleave(A_ ,							dim=0     )
        UpperCamelCase			:      Optional[Any] 			=     negative_image_embeds.repeat_interleave(A_ ,							dim=0     )

       UpperCamelCase			:      List[Any] 			=     torch.cat([negative_image_embeds, image_embeds] ,							dim=0     ).to(
           dtype=prompt_embeds.dtype ,							device=A_     )

       self.scheduler.set_timesteps(A_ ,							device=A_     )
       UpperCamelCase			:      Any 			=     self.scheduler.timesteps

       UpperCamelCase			:      Tuple 			=     self.unet.config.in_channels

       UpperCamelCase			:      Union[str, Any] 			=     get_new_h_w(A_ ,							A_ ,							self.movq_scale_factor     )

       # create initial latent
       UpperCamelCase			:      Tuple 			=     self.prepare_latents(
           (batch_size, num_channels_latents, height, width) ,							text_encoder_hidden_states.dtype ,							A_ ,							A_ ,							A_ ,							self.scheduler ,							)

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

        UpperCamelCase			:      Dict 			=     {"text_embeds": prompt_embeds, "image_embeds": image_embeds}
        UpperCamelCase			:      int 			=     self.unet(
            sample=A_ ,							timestep=A_ ,							encoder_hidden_states=A_ ,							added_cond_kwargs=A_ ,							return_dict=A_ ,							)[0]

        if do_classifier_free_guidance:
         UpperCamelCase			:      Dict 			=     noise_pred.split(latents.shape[1] ,							dim=1     )
         UpperCamelCase			:      str 			=     noise_pred.chunk(2     )
         UpperCamelCase			:      Optional[int] 			=     variance_pred.chunk(2     )
         UpperCamelCase			:      Dict 			=     noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
         UpperCamelCase			:      Any 			=     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"]
        ):
         UpperCamelCase			:      Dict 			=     noise_pred.split(latents.shape[1] ,							dim=1     )

        # compute the previous noisy sample x_t -> x_t-1
        UpperCamelCase			:      Optional[int] 			=     self.scheduler.step(
            A_ ,							A_ ,							A_ ,							generator=A_ ,							).prev_sample
       # post-processing
       UpperCamelCase			:      Tuple 			=     self.movq.decode(A_ ,							force_not_quantize=A_     )["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"]:
        UpperCamelCase			:      Tuple 			=     image * 0.5 + 0.5
        UpperCamelCase			:      Union[str, Any] 			=     image.clamp(0 ,							1     )
        UpperCamelCase			:      str 			=     image.cpu().permute(0 ,							2 ,							3 ,							1     ).float().numpy()

       if output_type == "pil":
        UpperCamelCase			:      Tuple 			=     self.numpy_to_pil(A_     )

       if not return_dict:
        return (image,)

       return ImagePipelineOutput(images=A_     )







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 	A_       (  _lowerCAmelCase  )		->				Union[str, Any]:  # picklable for multiprocessing
 return x.sum()



def 	A_       (  _lowerCAmelCase  )		->				Optional[Any]:  # picklable for multiprocessing
 return i + 1



@dataclass
class       A__     :
      _UpperCAmelCase			:int
      _UpperCAmelCase			:str



class       A__     ( __snake_case    ):





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Optional[int] 			=     {}
       UpperCamelCase			:      Optional[Any] 			=     []
       UpperCamelCase			:      List[Any] 			=     1
       UpperCamelCase			:      Tuple 			=     [1, 2]
       UpperCamelCase			:      Optional[Any] 			=     {"a": 1, "b": 2}
       UpperCamelCase			:      Optional[Any] 			=     {"a": [1, 2], "b": [3, 4]}
       UpperCamelCase			:      Any 			=     {"a": {"1": 1}, "b": 2}
       UpperCamelCase			:      List[str] 			=     {"a": 1, "b": 2, "c": 3, "d": 4}
       UpperCamelCase			:      Dict 			=     {}
       UpperCamelCase			:      Any 			=     []
       UpperCamelCase			:      Any 			=     2
       UpperCamelCase			:      Any 			=     [2, 3]
       UpperCamelCase			:      Optional[Any] 			=     {"a": 2, "b": 3}
       UpperCamelCase			:      List[Any] 			=     {"a": [2, 3], "b": [4, 5]}
       UpperCamelCase			:      Tuple 			=     {"a": {"1": 2}, "b": 3}
       UpperCamelCase			:      Dict 			=     {"a": 2, "b": 3, "c": 4, "d": 5}
       self.assertEqual(map_nested(A_ ,							A_     ) ,							A_     )
       self.assertEqual(map_nested(A_ ,							A_     ) ,							A_     )
       self.assertEqual(map_nested(A_ ,							A_     ) ,							A_     )
       self.assertEqual(map_nested(A_ ,							A_     ) ,							A_     )
       self.assertEqual(map_nested(A_ ,							A_     ) ,							A_     )
       self.assertEqual(map_nested(A_ ,							A_     ) ,							A_     )
       self.assertEqual(map_nested(A_ ,							A_     ) ,							A_     )
       self.assertEqual(map_nested(A_ ,							A_     ) ,							A_     )

       UpperCamelCase			:      List[str] 			=     2
       self.assertEqual(map_nested(A_ ,							A_ ,							num_proc=A_     ) ,							A_     )
       self.assertEqual(map_nested(A_ ,							A_ ,							num_proc=A_     ) ,							A_     )
       self.assertEqual(map_nested(A_ ,							A_ ,							num_proc=A_     ) ,							A_     )
       self.assertEqual(map_nested(A_ ,							A_ ,							num_proc=A_     ) ,							A_     )
       self.assertEqual(map_nested(A_ ,							A_ ,							num_proc=A_     ) ,							A_     )
       self.assertEqual(map_nested(A_ ,							A_ ,							num_proc=A_     ) ,							A_     )
       self.assertEqual(map_nested(A_ ,							A_ ,							num_proc=A_     ) ,							A_     )
       self.assertEqual(map_nested(A_ ,							A_ ,							num_proc=A_     ) ,							A_     )

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





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Optional[Any] 			=     {"a": 1, "b": 2}
       UpperCamelCase			:      List[Any] 			=     {"a": 3, "b": 4}
       UpperCamelCase			:      Tuple 			=     {"a": 5, "b": 6}
       UpperCamelCase			:      Union[str, Any] 			=     sorted([("a", (1, 3, 5)), ("b", (2, 4, 6))]     )
       self.assertEqual(sorted(zip_dict(A_ ,							A_ ,							A_     )     ) ,							A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       class       A__     :
             _UpperCAmelCase			:str							=     'bar'


       UpperCamelCase			:      List[Any] 			=     Foo()
       self.assertEqual(foo.my_attr ,							"bar"     )
       with temporary_assignment(A_ ,							"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 	A_       (  _lowerCAmelCase							,     _lowerCAmelCase							,     _lowerCAmelCase  )		->				Optional[Any]:
 with patch("datasets.utils.py_utils._single_map_nested"  ) as mock_single_map_nested, patch(
     "datasets.parallel.parallel.Pool"  ) as mock_multiprocessing_pool:
  UpperCamelCase			:      Union[str, Any] 			=     {F"""{i}""": i for i in range(_lowerCAmelCase  )}
  UpperCamelCase			:      List[str] 			=     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__     ( __snake_case    ):





      @require_tf
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       import tensorflow as tf
       from tensorflow.keras import layers

       UpperCamelCase			:      int 			=     layers.Dense(2     )

       def gen_random_output():
        UpperCamelCase			:      Optional[Any] 			=     tf.random.uniform((1, 3)     )
        return model(A_     ).numpy()

       with temp_seed(42 ,							set_tensorflow=A_     ):
        UpperCamelCase			:      List[Any] 			=     gen_random_output()
       with temp_seed(42 ,							set_tensorflow=A_     ):
        UpperCamelCase			:      Dict 			=     gen_random_output()
       UpperCamelCase			:      Optional[int] 			=     gen_random_output()

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





      @require_torch
      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       import torch

       def gen_random_output():
        UpperCamelCase			:      Optional[Any] 			=     torch.nn.Linear(3 ,							2     )
        UpperCamelCase			:      Dict 			=     torch.rand(1 ,							3     )
        return model(A_     ).detach().numpy()

       with temp_seed(42 ,							set_pytorch=A_     ):
        UpperCamelCase			:      Dict 			=     gen_random_output()
       with temp_seed(42 ,							set_pytorch=A_     ):
        UpperCamelCase			:      Optional[int] 			=     gen_random_output()
       UpperCamelCase			:      List[Any] 			=     gen_random_output()

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





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       def gen_random_output():
        return np.random.rand(1 ,							3     )

       with temp_seed(42     ):
        UpperCamelCase			:      Optional[Any] 			=     gen_random_output()
       with temp_seed(42     ):
        UpperCamelCase			:      Optional[Any] 			=     gen_random_output()
       UpperCamelCase			:      Optional[Any] 			=     gen_random_output()

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



@pytest.mark.parametrize("input_data"							,     [{}]  )
def 	A_       (  _lowerCAmelCase  )		->				List[Any]:
 UpperCamelCase			:      Optional[Any] 			=     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 	A_       (  _lowerCAmelCase							,     _lowerCAmelCase  )		->				Tuple:
 UpperCamelCase			:      Dict 			=     NestedDataStructure(_lowerCAmelCase  ).flatten()
 assert output == expected_output



def 	A_       (  )		->				List[Any]:
 UpperCamelCase			:      str 			=     A(x=1							,     y="foobar"  )
 UpperCamelCase			:      Tuple 			=     {"x": 1, "y": "foobar"}
 assert asdict(_lowerCAmelCase  ) == expected_output

 UpperCamelCase			:      List[str] 			=     {"a": {"b": A(x=10							,     y="foo"  )}, "c": [A(x=20							,     y="bar"  )]}
 UpperCamelCase			:      Tuple 			=     {"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 	A_       (  _lowerCAmelCase  )		->				Tuple:
 return text.split()



def 	A_       (  _lowerCAmelCase  )		->				Dict:
 yield (time.time(), content)
 time.sleep(2  )
 yield (time.time(), content)



def 	A_       (  )		->				str:
 with Pool(2  ) as pool:
  UpperCamelCase			:      List[str] 			=     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:
  UpperCamelCase			:      Dict 			=     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:
  UpperCamelCase			:      Any 			=     []
  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
