import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotSmallConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
    import os

    # The slow tests are often failing with OOM error on GPU
    # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
    # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
    __lowercase											=      '''platform'''

    import jax
    import jax.numpy as jnp

    from transformers.models.blenderbot_small.modeling_flax_blenderbot_small import (
        FlaxBlenderbotSmallForConditionalGeneration,
        FlaxBlenderbotSmallModel,
        shift_tokens_right,
    )



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE=None						,    SCREAMING_SNAKE_CASE=None						,    SCREAMING_SNAKE_CASE=None						,    SCREAMING_SNAKE_CASE=None						,    SCREAMING_SNAKE_CASE=None						,    SCREAMING_SNAKE_CASE=None						,    ):
    '''simple docstring'''







    if attention_mask is None:
        __UpperCamelCase  :int    						=	np.where(input_ids != config.pad_token_id						,    1						,    0			)
    if decoder_attention_mask is None:
        __UpperCamelCase  :Tuple    						=	np.where(decoder_input_ids != config.pad_token_id						,    1						,    0			)
    if head_mask is None:
        __UpperCamelCase  :Union[str, Any]    						=	np.ones((config.encoder_layers, config.encoder_attention_heads)			)
    if decoder_head_mask is None:
        __UpperCamelCase  :Optional[int]    						=	np.ones((config.decoder_layers, config.decoder_attention_heads)			)
    if cross_attn_head_mask is None:
        __UpperCamelCase  :int    						=	np.ones((config.decoder_layers, config.decoder_attention_heads)			)
    return {
        "input_ids": input_ids,
        "decoder_input_ids": decoder_input_ids,
        "attention_mask": attention_mask,
        "decoder_attention_mask": attention_mask,
    }


class 			lowerCamelCase_  :






    '''simple docstring'''






    def __init__(			self				,				__lowercase				,				__lowercase=13				,				__lowercase=7				,				__lowercase=True				,				__lowercase=False				,				__lowercase=99				,				__lowercase=16				,				__lowercase=2				,				__lowercase=4				,				__lowercase=4				,				__lowercase="gelu"				,				__lowercase=0.1				,				__lowercase=0.1				,				__lowercase=32				,				__lowercase=2				,				__lowercase=1				,				__lowercase=0				,				__lowercase=0.02				,				) ->			List[str]:
        __UpperCamelCase  :Union[str, Any]    						=	parent
        __UpperCamelCase  :List[str]    						=	batch_size
        __UpperCamelCase  :Any    						=	seq_length
        __UpperCamelCase  :Tuple    						=	is_training
        __UpperCamelCase  :Optional[Any]    						=	use_labels
        __UpperCamelCase  :int    						=	vocab_size
        __UpperCamelCase  :int    						=	hidden_size
        __UpperCamelCase  :Optional[Any]    						=	num_hidden_layers
        __UpperCamelCase  :Any    						=	num_attention_heads
        __UpperCamelCase  :Optional[int]    						=	intermediate_size
        __UpperCamelCase  :Tuple    						=	hidden_act
        __UpperCamelCase  :List[str]    						=	hidden_dropout_prob
        __UpperCamelCase  :Optional[Any]    						=	attention_probs_dropout_prob
        __UpperCamelCase  :List[Any]    						=	max_position_embeddings
        __UpperCamelCase  :str    						=	eos_token_id
        __UpperCamelCase  :Any    						=	pad_token_id
        __UpperCamelCase  :str    						=	bos_token_id
        __UpperCamelCase  :Union[str, Any]    						=	initializer_range







    def     UpperCamelCase__     (			self) ->			str:
        __UpperCamelCase  :List[str]    						=	np.clip(ids_tensor([self.batch_size, self.seq_length - 1]				,				self.vocab_size)				,				3				,				self.vocab_size)
        __UpperCamelCase  :Optional[int]    						=	np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)				,				dtype=np.intaa))				,				-1)

        __UpperCamelCase  :str    						=	shift_tokens_right(__lowercase				,				1				,				2)

        __UpperCamelCase  :str    						=	BlenderbotSmallConfig(
            vocab_size=self.vocab_size				,				d_model=self.hidden_size				,				encoder_layers=self.num_hidden_layers				,				decoder_layers=self.num_hidden_layers				,				encoder_attention_heads=self.num_attention_heads				,				decoder_attention_heads=self.num_attention_heads				,				encoder_ffn_dim=self.intermediate_size				,				decoder_ffn_dim=self.intermediate_size				,				dropout=self.hidden_dropout_prob				,				attention_dropout=self.attention_probs_dropout_prob				,				max_position_embeddings=self.max_position_embeddings				,				eos_token_id=self.eos_token_id				,				bos_token_id=self.bos_token_id				,				pad_token_id=self.pad_token_id				,				initializer_range=self.initializer_range				,				use_cache=__lowercase				,				)
        __UpperCamelCase  :List[Any]    						=	prepare_blenderbot_inputs_dict(__lowercase				,				__lowercase				,				__lowercase)
        return config, inputs_dict







    def     UpperCamelCase__     (			self) ->			Dict:
        __UpperCamelCase						, __UpperCamelCase  :Optional[Any]    						=	self.prepare_config_and_inputs()
        return config, inputs_dict







    def     UpperCamelCase__     (			self				,				__lowercase				,				__lowercase				,				__lowercase) ->			List[str]:
        __UpperCamelCase  :Dict    						=	20
        __UpperCamelCase  :str    						=	model_class_name(__lowercase)

        __UpperCamelCase  :Dict    						=	model.encode(inputs_dict['''input_ids'''])

        __UpperCamelCase						, __UpperCamelCase  :Dict    						=	(
            inputs_dict['''decoder_input_ids'''],
            inputs_dict['''decoder_attention_mask'''],
        )

        __UpperCamelCase  :List[str]    						=	model.init_cache(decoder_input_ids.shape[0]				,				__lowercase				,				__lowercase)
        __UpperCamelCase  :Optional[Any]    						=	jnp.ones((decoder_input_ids.shape[0], max_decoder_length)				,				dtype='''i4''')

        __UpperCamelCase  :Union[str, Any]    						=	jnp.broadcast_to(
            jnp.arange(decoder_input_ids.shape[-1] - 1)[None, :]				,				(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)				,				)
        __UpperCamelCase  :Optional[Any]    						=	model.decode(
            decoder_input_ids[:, :-1]				,				__lowercase				,				decoder_attention_mask=__lowercase				,				past_key_values=__lowercase				,				decoder_position_ids=__lowercase				,				)

        __UpperCamelCase  :str    						=	jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]				,				dtype='''i4''')
        __UpperCamelCase  :Optional[int]    						=	model.decode(
            decoder_input_ids[:, -1:]				,				__lowercase				,				decoder_attention_mask=__lowercase				,				past_key_values=outputs_cache.past_key_values				,				decoder_position_ids=__lowercase				,				)

        __UpperCamelCase  :List[Any]    						=	model.decode(__lowercase				,				__lowercase)

        __UpperCamelCase  :Dict    						=	np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])))
        self.parent.assertTrue(diff < 1E-3				,				msg=f"""Max diff is {diff}""")







    def     UpperCamelCase__     (			self				,				__lowercase				,				__lowercase				,				__lowercase) ->			str:
        __UpperCamelCase  :Optional[Any]    						=	20
        __UpperCamelCase  :Dict    						=	model_class_name(__lowercase)

        __UpperCamelCase  :Optional[Any]    						=	model.encode(inputs_dict['''input_ids'''])

        __UpperCamelCase						, __UpperCamelCase  :List[Any]    						=	(
            inputs_dict['''decoder_input_ids'''],
            inputs_dict['''decoder_attention_mask'''],
        )

        __UpperCamelCase  :str    						=	jnp.concatenate(
            [
                decoder_attention_mask,
                jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])),
            ]				,				axis=-1				,				)

        __UpperCamelCase  :Union[str, Any]    						=	model.init_cache(decoder_input_ids.shape[0]				,				__lowercase				,				__lowercase)
        __UpperCamelCase  :Optional[Any]    						=	jnp.broadcast_to(
            jnp.arange(decoder_input_ids.shape[-1] - 1)[None, :]				,				(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)				,				)

        __UpperCamelCase  :Dict    						=	model.decode(
            decoder_input_ids[:, :-1]				,				__lowercase				,				decoder_attention_mask=__lowercase				,				past_key_values=__lowercase				,				decoder_position_ids=__lowercase				,				)
        __UpperCamelCase  :Union[str, Any]    						=	jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]				,				dtype='''i4''')
        __UpperCamelCase  :List[Any]    						=	model.decode(
            decoder_input_ids[:, -1:]				,				__lowercase				,				past_key_values=outputs_cache.past_key_values				,				decoder_attention_mask=__lowercase				,				decoder_position_ids=__lowercase				,				)

        __UpperCamelCase  :int    						=	model.decode(__lowercase				,				__lowercase				,				decoder_attention_mask=__lowercase)

        __UpperCamelCase  :Any    						=	np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])))
        self.parent.assertTrue(diff < 1E-3				,				msg=f"""Max diff is {diff}""")


@require_flax
class 			lowerCamelCase_  (     unittest.TestCase				):






    '''simple docstring'''






    a__			:		Any      			=   9_9







    def     UpperCamelCase__     (			self) ->			Optional[int]:
        __UpperCamelCase  :List[str]    						=	np.array(
            [
                [71, 82, 18, 33, 46, 91, 2],
                [68, 34, 26, 58, 30, 82, 2],
                [5, 97, 17, 39, 94, 40, 2],
                [76, 83, 94, 25, 70, 78, 2],
                [87, 59, 41, 35, 48, 66, 2],
                [55, 13, 16, 58, 5, 2, 1],  # note padding
                [64, 27, 31, 51, 12, 75, 2],
                [52, 64, 86, 17, 83, 39, 2],
                [48, 61, 9, 24, 71, 82, 2],
                [26, 1, 60, 48, 22, 13, 2],
                [21, 5, 62, 28, 14, 76, 2],
                [45, 98, 37, 86, 59, 48, 2],
                [70, 70, 50, 9, 28, 0, 2],
            ]				,				dtype=np.intaa				,				)

        __UpperCamelCase  :Tuple    						=	input_ids.shape[0]
        __UpperCamelCase  :Tuple    						=	BlenderbotSmallConfig(
            vocab_size=self.vocab_size				,				d_model=24				,				encoder_layers=2				,				decoder_layers=2				,				encoder_attention_heads=2				,				decoder_attention_heads=2				,				encoder_ffn_dim=32				,				decoder_ffn_dim=32				,				max_position_embeddings=48				,				eos_token_id=2				,				pad_token_id=1				,				bos_token_id=0				,				)
        return config, input_ids, batch_size







    def     UpperCamelCase__     (			self) ->			int:
        __UpperCamelCase						, __UpperCamelCase						, __UpperCamelCase  :str    						=	self._get_config_and_data()
        __UpperCamelCase  :str    						=	FlaxBlenderbotSmallForConditionalGeneration(__lowercase)
        __UpperCamelCase  :List[Any]    						=	lm_model(input_ids=__lowercase)
        __UpperCamelCase  :List[str]    						=	(batch_size, input_ids.shape[1], config.vocab_size)
        self.assertEqual(outputs['''logits'''].shape				,				__lowercase)







    def     UpperCamelCase__     (			self) ->			Dict:
        __UpperCamelCase  :int    						=	BlenderbotSmallConfig(
            vocab_size=self.vocab_size				,				d_model=14				,				encoder_layers=2				,				decoder_layers=2				,				encoder_attention_heads=2				,				decoder_attention_heads=2				,				encoder_ffn_dim=8				,				decoder_ffn_dim=8				,				max_position_embeddings=48				,				)
        __UpperCamelCase  :Optional[int]    						=	FlaxBlenderbotSmallForConditionalGeneration(__lowercase)
        __UpperCamelCase  :Optional[int]    						=	np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]				,				dtype=np.intaa)
        __UpperCamelCase  :Union[str, Any]    						=	np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]				,				dtype=np.intaa)
        __UpperCamelCase  :str    						=	lm_model(input_ids=__lowercase				,				decoder_input_ids=__lowercase)
        __UpperCamelCase  :List[Any]    						=	(*summary.shape, config.vocab_size)
        self.assertEqual(outputs['''logits'''].shape				,				__lowercase)







    def     UpperCamelCase__     (			self) ->			Optional[int]:
        __UpperCamelCase  :List[Any]    						=	np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]				,				dtype=np.intaa)
        __UpperCamelCase  :str    						=	shift_tokens_right(__lowercase				,				1				,				2)
        __UpperCamelCase  :Dict    						=	np.equal(__lowercase				,				1).astype(np.floataa).sum()
        __UpperCamelCase  :Optional[int]    						=	np.equal(__lowercase				,				1).astype(np.floataa).sum()
        self.assertEqual(shifted.shape				,				input_ids.shape)
        self.assertEqual(__lowercase				,				n_pad_before - 1)
        self.assertTrue(np.equal(shifted[:, 0]				,				2).all())


@require_flax
class 			lowerCamelCase_  (     UpperCAmelCase_		,			unittest.TestCase		,			UpperCAmelCase_				):






    '''simple docstring'''






    a__			:		Optional[Any]      			=   True
    a__			:		int      			=   (
        (
            FlaxBlenderbotSmallModel,
            FlaxBlenderbotSmallForConditionalGeneration,
        )
        if is_flax_available()
        else ()
    )
    a__			:		Union[str, Any]      			=   (FlaxBlenderbotSmallForConditionalGeneration,) if is_flax_available() else ()







    def     UpperCamelCase__     (			self) ->			Tuple:
        __UpperCamelCase  :str    						=	FlaxBlenderbotSmallModelTester(self)







    def     UpperCamelCase__     (			self) ->			Dict:
        __UpperCamelCase						, __UpperCamelCase  :List[Any]    						=	self.model_tester.prepare_config_and_inputs()
        for model_class in self.all_model_classes:
            self.model_tester.check_use_cache_forward(__lowercase				,				__lowercase				,				__lowercase)







    def     UpperCamelCase__     (			self) ->			Union[str, Any]:
        __UpperCamelCase						, __UpperCamelCase  :Optional[int]    						=	self.model_tester.prepare_config_and_inputs()
        for model_class in self.all_model_classes:
            self.model_tester.check_use_cache_forward_with_attn_mask(__lowercase				,				__lowercase				,				__lowercase)







    def     UpperCamelCase__     (			self) ->			List[str]:
        __UpperCamelCase						, __UpperCamelCase  :List[str]    						=	self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            with self.subTest(model_class.__name__):
                __UpperCamelCase  :List[str]    						=	self._prepare_for_class(__lowercase				,				__lowercase)
                __UpperCamelCase  :str    						=	model_class(__lowercase)

                @jax.jit
                def encode_jitted(__lowercase				,				__lowercase=None				,				**__lowercase):
                    return model.encode(input_ids=__lowercase				,				attention_mask=__lowercase)

                with self.subTest('''JIT Enabled'''):
                    __UpperCamelCase  :List[str]    						=	encode_jitted(**__lowercase).to_tuple()

                with self.subTest('''JIT Disabled'''):
                    with jax.disable_jit():
                        __UpperCamelCase  :str    						=	encode_jitted(**__lowercase).to_tuple()

                self.assertEqual(len(__lowercase)				,				len(__lowercase))
                for jitted_output, output in zip(__lowercase				,				__lowercase):
                    self.assertEqual(jitted_output.shape				,				output.shape)







    def     UpperCamelCase__     (			self) ->			Tuple:
        __UpperCamelCase						, __UpperCamelCase  :Dict    						=	self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            with self.subTest(model_class.__name__):
                __UpperCamelCase  :Union[str, Any]    						=	model_class(__lowercase)
                __UpperCamelCase  :Dict    						=	model.encode(inputs_dict['''input_ids''']				,				inputs_dict['''attention_mask'''])

                __UpperCamelCase  :Tuple    						=	{
                    '''decoder_input_ids''': inputs_dict['''decoder_input_ids'''],
                    '''decoder_attention_mask''': inputs_dict['''decoder_attention_mask'''],
                    '''encoder_outputs''': encoder_outputs,
                }

                @jax.jit
                def decode_jitted(__lowercase				,				__lowercase				,				__lowercase):
                    return model.decode(
                        decoder_input_ids=__lowercase				,				decoder_attention_mask=__lowercase				,				encoder_outputs=__lowercase				,				)

                with self.subTest('''JIT Enabled'''):
                    __UpperCamelCase  :Union[str, Any]    						=	decode_jitted(**__lowercase).to_tuple()

                with self.subTest('''JIT Disabled'''):
                    with jax.disable_jit():
                        __UpperCamelCase  :Union[str, Any]    						=	decode_jitted(**__lowercase).to_tuple()

                self.assertEqual(len(__lowercase)				,				len(__lowercase))
                for jitted_output, output in zip(__lowercase				,				__lowercase):
                    self.assertEqual(jitted_output.shape				,				output.shape)







    @slow
    def     UpperCamelCase__     (			self) ->			List[Any]:
        for model_class_name in self.all_model_classes:
            __UpperCamelCase  :Tuple    						=	model_class_name.from_pretrained('''facebook/blenderbot_small-90M''')
            # FlaxBlenderbotForSequenceClassification expects eos token in input_ids
            __UpperCamelCase  :Tuple    						=	np.ones((1, 1)) * model.config.eos_token_id
            __UpperCamelCase  :List[Any]    						=	model(__lowercase)
            self.assertIsNotNone(__lowercase)






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						(     SCREAMING_SNAKE_CASE			):  # picklable for multiprocessing
    '''simple docstring'''







    return x.sum()



def lowerCamelCase						(     SCREAMING_SNAKE_CASE			):  # picklable for multiprocessing
    '''simple docstring'''







    return i + 1


@dataclass
class 			lowerCamelCase_  :






    '''simple docstring'''






    a__			:		int
    a__			:		str


class 			lowerCamelCase_  (     UpperCAmelCase_				):






    '''simple docstring'''






    def     UpperCamelCase__     (			self) ->			int:
        __UpperCamelCase  :Dict    						=	{}
        __UpperCamelCase  :int    						=	[]
        __UpperCamelCase  :int    						=	1
        __UpperCamelCase  :int    						=	[1, 2]
        __UpperCamelCase  :Optional[Any]    						=	{'''a''': 1, '''b''': 2}
        __UpperCamelCase  :Dict    						=	{'''a''': [1, 2], '''b''': [3, 4]}
        __UpperCamelCase  :Optional[int]    						=	{'''a''': {'''1''': 1}, '''b''': 2}
        __UpperCamelCase  :Optional[Any]    						=	{'''a''': 1, '''b''': 2, '''c''': 3, '''d''': 4}
        __UpperCamelCase  :str    						=	{}
        __UpperCamelCase  :Any    						=	[]
        __UpperCamelCase  :Optional[Any]    						=	2
        __UpperCamelCase  :Optional[int]    						=	[2, 3]
        __UpperCamelCase  :List[Any]    						=	{'''a''': 2, '''b''': 3}
        __UpperCamelCase  :List[str]    						=	{'''a''': [2, 3], '''b''': [4, 5]}
        __UpperCamelCase  :int    						=	{'''a''': {'''1''': 2}, '''b''': 3}
        __UpperCamelCase  :Union[str, Any]    						=	{'''a''': 2, '''b''': 3, '''c''': 4, '''d''': 5}
        self.assertEqual(map_nested(__lowercase				,				__lowercase)				,				__lowercase)
        self.assertEqual(map_nested(__lowercase				,				__lowercase)				,				__lowercase)
        self.assertEqual(map_nested(__lowercase				,				__lowercase)				,				__lowercase)
        self.assertEqual(map_nested(__lowercase				,				__lowercase)				,				__lowercase)
        self.assertEqual(map_nested(__lowercase				,				__lowercase)				,				__lowercase)
        self.assertEqual(map_nested(__lowercase				,				__lowercase)				,				__lowercase)
        self.assertEqual(map_nested(__lowercase				,				__lowercase)				,				__lowercase)
        self.assertEqual(map_nested(__lowercase				,				__lowercase)				,				__lowercase)

        __UpperCamelCase  :Optional[Any]    						=	2
        self.assertEqual(map_nested(__lowercase				,				__lowercase				,				num_proc=__lowercase)				,				__lowercase)
        self.assertEqual(map_nested(__lowercase				,				__lowercase				,				num_proc=__lowercase)				,				__lowercase)
        self.assertEqual(map_nested(__lowercase				,				__lowercase				,				num_proc=__lowercase)				,				__lowercase)
        self.assertEqual(map_nested(__lowercase				,				__lowercase				,				num_proc=__lowercase)				,				__lowercase)
        self.assertEqual(map_nested(__lowercase				,				__lowercase				,				num_proc=__lowercase)				,				__lowercase)
        self.assertEqual(map_nested(__lowercase				,				__lowercase				,				num_proc=__lowercase)				,				__lowercase)
        self.assertEqual(map_nested(__lowercase				,				__lowercase				,				num_proc=__lowercase)				,				__lowercase)
        self.assertEqual(map_nested(__lowercase				,				__lowercase				,				num_proc=__lowercase)				,				__lowercase)

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







    def     UpperCamelCase__     (			self) ->			List[Any]:
        __UpperCamelCase  :List[str]    						=	{'''a''': 1, '''b''': 2}
        __UpperCamelCase  :int    						=	{'''a''': 3, '''b''': 4}
        __UpperCamelCase  :List[str]    						=	{'''a''': 5, '''b''': 6}
        __UpperCamelCase  :List[Any]    						=	sorted([('''a''', (1, 3, 5)), ('''b''', (2, 4, 6))])
        self.assertEqual(sorted(zip_dict(__lowercase				,				__lowercase				,				__lowercase))				,				__lowercase)







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


        class 			lowerCamelCase_  :






            '''simple docstring'''






            a__			:		int      			=   """bar"""
        __UpperCamelCase  :List[str]    						=	Foo()
        self.assertEqual(foo.my_attr				,				'''bar''')
        with temporary_assignment(__lowercase				,				'''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						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    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(SCREAMING_SNAKE_CASE			)}
        __UpperCamelCase  :Dict    						=	map_nested(lambda SCREAMING_SNAKE_CASE			: x + 10						,    SCREAMING_SNAKE_CASE						,    num_proc=SCREAMING_SNAKE_CASE						,    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 			lowerCamelCase_  (     UpperCAmelCase_				):






    '''simple docstring'''






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

        __UpperCamelCase  :Optional[Any]    						=	layers.Dense(2)

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

        with temp_seed(42				,				set_tensorflow=__lowercase):
            __UpperCamelCase  :Tuple    						=	gen_random_output()
        with temp_seed(42				,				set_tensorflow=__lowercase):
            __UpperCamelCase  :int    						=	gen_random_output()
        __UpperCamelCase  :List[str]    						=	gen_random_output()

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







    @require_torch
    def     UpperCamelCase__     (			self) ->			List[str]:
        import torch

        def gen_random_output():
            __UpperCamelCase  :List[Any]    						=	torch.nn.Linear(3				,				2)
            __UpperCamelCase  :Optional[Any]    						=	torch.rand(1				,				3)
            return model(__lowercase).detach().numpy()

        with temp_seed(42				,				set_pytorch=__lowercase):
            __UpperCamelCase  :Union[str, Any]    						=	gen_random_output()
        with temp_seed(42				,				set_pytorch=__lowercase):
            __UpperCamelCase  :List[str]    						=	gen_random_output()
        __UpperCamelCase  :Dict    						=	gen_random_output()

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







    def     UpperCamelCase__     (			self) ->			str:
        def gen_random_output():
            return np.random.rand(1				,				3)

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

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



@pytest.mark.parametrize('''input_data'''						,    [{}]			)
def lowerCamelCase						(     SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    __UpperCamelCase  :List[str]    						=	NestedDataStructure(SCREAMING_SNAKE_CASE			).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						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    __UpperCamelCase  :Dict    						=	NestedDataStructure(SCREAMING_SNAKE_CASE			).flatten()
    assert output == expected_output



def lowerCamelCase						(     ):
    '''simple docstring'''







    __UpperCamelCase  :Union[str, Any]    						=	A(x=1						,    y='''foobar'''			)
    __UpperCamelCase  :Dict    						=	{'''x''': 1, '''y''': '''foobar'''}
    assert asdict(SCREAMING_SNAKE_CASE			) == expected_output

    __UpperCamelCase  :str    						=	{'''a''': {'''b''': A(x=10						,    y='''foo'''			)}, '''c''': [A(x=20						,    y='''bar'''			)]}
    __UpperCamelCase  :List[Any]    						=	{'''a''': {'''b''': {'''x''': 10, '''y''': '''foo'''}}, '''c''': [{'''x''': 20, '''y''': '''bar'''}]}
    assert asdict(SCREAMING_SNAKE_CASE			) == expected_output

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



def lowerCamelCase						(     SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    return text.split()



def lowerCamelCase						(     SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    yield (time.time(), content)
    time.sleep(2			)
    yield (time.time(), content)



def lowerCamelCase						(     ):
    '''simple docstring'''







    with Pool(2			) as pool:
        __UpperCamelCase  :Tuple    						=	list(iflatmap_unordered(SCREAMING_SNAKE_CASE						,    _split_text						,    kwargs_iterable=[{'''text''': '''hello there'''}] * 10			)			)
        assert out.count('''hello'''			) == 10
        assert out.count('''there'''			) == 10
        assert len(SCREAMING_SNAKE_CASE			) == 20

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

    # check that we get items as fast as possible
    with Pool(2			) as pool:
        __UpperCamelCase  :Dict    						=	[]
        for yield_time, content in iflatmap_unordered(
            SCREAMING_SNAKE_CASE						,    _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(SCREAMING_SNAKE_CASE			)
        assert out.count('''a'''			) == 2
        assert out.count('''b'''			) == 2
        assert len(SCREAMING_SNAKE_CASE			) == 4





