
import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
						import jax
						import jax.numpy as jnp

						from transformers.generation import (
						    FlaxForcedBOSTokenLogitsProcessor,
						    FlaxForcedEOSTokenLogitsProcessor,
						    FlaxLogitsProcessorList,
						    FlaxMinLengthLogitsProcessor,
						    FlaxTemperatureLogitsWarper,
						    FlaxTopKLogitsWarper,
						    FlaxTopPLogitsWarper,
						)





@require_flax
class  snake_case     (							unittest.TestCase   ):

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = jnp.ones((batch_size, length)    ) / length
									return scores

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = None
									SCREAMING_SNAKE_CASE_						 = 20

									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(batch_size=2       ,       length=lowercase__    )

									# tweak scores to not be uniform anymore
									SCREAMING_SNAKE_CASE_						 = scores.at[1, 5].set((1 / length) + 0.1    )  # peak, 1st batch
									SCREAMING_SNAKE_CASE_						 = scores.at[1, 10].set((1 / length) - 0.4    )  # valley, 1st batch

									# compute softmax
									SCREAMING_SNAKE_CASE_						 = jax.nn.softmax(lowercase__       ,       axis=-1    )

									SCREAMING_SNAKE_CASE_						 = FlaxTemperatureLogitsWarper(temperature=0.5    )
									SCREAMING_SNAKE_CASE_						 = FlaxTemperatureLogitsWarper(temperature=1.3    )

									SCREAMING_SNAKE_CASE_						 = jax.nn.softmax(temp_dist_warper_sharper(lowercase__       ,       scores.copy()       ,       cur_len=lowercase__    )       ,       axis=-1    )
									SCREAMING_SNAKE_CASE_						 = jax.nn.softmax(temp_dist_warper_smoother(lowercase__       ,       scores.copy()       ,       cur_len=lowercase__    )       ,       axis=-1    )

									# uniform distribution stays uniform
									self.assertTrue(jnp.allclose(probs[0, :]       ,       warped_prob_sharp[0, :]       ,       atol=1e-3    )    )
									self.assertTrue(jnp.allclose(probs[0, :]       ,       warped_prob_smooth[0, :]       ,       atol=1e-3    )    )

									# sharp peaks get higher, valleys get lower
									self.assertLess(probs[1, :].max()       ,       warped_prob_sharp[1, :].max()    )
									self.assertGreater(probs[1, :].min()       ,       warped_prob_sharp[1, :].min()    )

									# smooth peaks get lower, valleys get higher
									self.assertGreater(probs[1, :].max()       ,       warped_prob_smooth[1, :].max()    )
									self.assertLess(probs[1, :].min()       ,       warped_prob_smooth[1, :].min()    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = None
									SCREAMING_SNAKE_CASE_						 = 10
									SCREAMING_SNAKE_CASE_						 = 2

									# create ramp distribution
									SCREAMING_SNAKE_CASE_						 = np.broadcast_to(np.arange(lowercase__    )[None, :]       ,       (batch_size, vocab_size)    ).copy()
									SCREAMING_SNAKE_CASE_						 = ramp_logits[1:, : vocab_size // 2] + vocab_size

									SCREAMING_SNAKE_CASE_						 = FlaxTopKLogitsWarper(3    )

									SCREAMING_SNAKE_CASE_						 = top_k_warp(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )

									# check that correct tokens are filtered
									self.assertListEqual(jnp.isinf(scores[0]    ).tolist()       ,       7 * [True] + 3 * [False]    )
									self.assertListEqual(jnp.isinf(scores[1]    ).tolist()       ,       2 * [True] + 3 * [False] + 5 * [True]    )

									# check special case
									SCREAMING_SNAKE_CASE_						 = 5
									SCREAMING_SNAKE_CASE_						 = FlaxTopKLogitsWarper(top_k=1       ,       filter_value=0.0       ,       min_tokens_to_keep=3    )

									SCREAMING_SNAKE_CASE_						 = np.broadcast_to(np.arange(lowercase__    )[None, :]       ,       (batch_size, length)    ).copy()
									SCREAMING_SNAKE_CASE_						 = top_k_warp_safety_check(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )

									# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
									self.assertListEqual((scores == 0.0).sum(axis=-1    ).tolist()       ,       [2, 2]    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = None
									SCREAMING_SNAKE_CASE_						 = 10
									SCREAMING_SNAKE_CASE_						 = 2

									# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
									SCREAMING_SNAKE_CASE_						 = np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]    )    )

									SCREAMING_SNAKE_CASE_						 = FlaxTopPLogitsWarper(0.8    )
									SCREAMING_SNAKE_CASE_						 = np.exp(top_p_warp(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )    )

									# dist should be filtered to keep min num values so that sum is >= top_p
									# exp (-inf) => 0
									SCREAMING_SNAKE_CASE_						 = np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]    )
									self.assertTrue(np.allclose(lowercase__       ,       lowercase__       ,       atol=1e-3    )    )

									# check edge cases with negative and extreme logits
									SCREAMING_SNAKE_CASE_						 = np.broadcast_to(np.arange(lowercase__    )[None, :]       ,       (batch_size, vocab_size)    ).copy() - (
									    vocab_size // 2
									)

									# make ramp_logits more extreme
									SCREAMING_SNAKE_CASE_						 = ramp_logits[1] * 1_00.0

									# make sure at least 2 tokens are kept
									SCREAMING_SNAKE_CASE_						 = FlaxTopPLogitsWarper(0.9       ,       min_tokens_to_keep=2       ,       filter_value=0.0    )
									SCREAMING_SNAKE_CASE_						 = top_p_warp(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )

									# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
									self.assertListEqual((filtered_dist != 0.0).sum(axis=-1    ).tolist()       ,       [3, 2]    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = 20
									SCREAMING_SNAKE_CASE_						 = 4
									SCREAMING_SNAKE_CASE_						 = 0

									SCREAMING_SNAKE_CASE_						 = FlaxMinLengthLogitsProcessor(min_length=10       ,       eos_token_id=lowercase__    )

									# check that min length is applied at length 5
									SCREAMING_SNAKE_CASE_						 = ids_tensor((batch_size, 20)       ,       vocab_size=20    )
									SCREAMING_SNAKE_CASE_						 = 5
									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(lowercase__       ,       lowercase__    )
									SCREAMING_SNAKE_CASE_						 = min_dist_processor(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
									self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()       ,       4 * [-float('''inf'''    )]    )

									# check that min length is not applied anymore at length 15
									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(lowercase__       ,       lowercase__    )
									SCREAMING_SNAKE_CASE_						 = 15
									SCREAMING_SNAKE_CASE_						 = min_dist_processor(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
									self.assertFalse(jnp.isinf(lowercase__    ).any()    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = 20
									SCREAMING_SNAKE_CASE_						 = 4
									SCREAMING_SNAKE_CASE_						 = 0

									SCREAMING_SNAKE_CASE_						 = FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowercase__    )

									# check that all scores are -inf except the bos_token_id score
									SCREAMING_SNAKE_CASE_						 = ids_tensor((batch_size, 1)       ,       vocab_size=20    )
									SCREAMING_SNAKE_CASE_						 = 1
									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(lowercase__       ,       lowercase__    )
									SCREAMING_SNAKE_CASE_						 = logits_processor(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
									self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]    ).all()    )
									self.assertListEqual(scores[:, bos_token_id].tolist()       ,       4 * [0]    )  # score for bos_token_id shold be zero

									# check that bos_token_id is not forced if current length is greater than 1
									SCREAMING_SNAKE_CASE_						 = 3
									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(lowercase__       ,       lowercase__    )
									SCREAMING_SNAKE_CASE_						 = logits_processor(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
									self.assertFalse(jnp.isinf(lowercase__    ).any()    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = 20
									SCREAMING_SNAKE_CASE_						 = 4
									SCREAMING_SNAKE_CASE_						 = 0
									SCREAMING_SNAKE_CASE_						 = 5

									SCREAMING_SNAKE_CASE_						 = FlaxForcedEOSTokenLogitsProcessor(max_length=lowercase__       ,       eos_token_id=lowercase__    )

									# check that all scores are -inf except the eos_token_id when max_length is reached
									SCREAMING_SNAKE_CASE_						 = ids_tensor((batch_size, 4)       ,       vocab_size=20    )
									SCREAMING_SNAKE_CASE_						 = 4
									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(lowercase__       ,       lowercase__    )
									SCREAMING_SNAKE_CASE_						 = logits_processor(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
									self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]    ).all()    )
									self.assertListEqual(scores[:, eos_token_id].tolist()       ,       4 * [0]    )  # score for eos_token_id should be zero

									# check that eos_token_id is not forced if max_length is not reached
									SCREAMING_SNAKE_CASE_						 = 3
									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(lowercase__       ,       lowercase__    )
									SCREAMING_SNAKE_CASE_						 = logits_processor(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
									self.assertFalse(jnp.isinf(lowercase__    ).any()    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = 4
									SCREAMING_SNAKE_CASE_						 = 10
									SCREAMING_SNAKE_CASE_						 = 15
									SCREAMING_SNAKE_CASE_						 = 2
									SCREAMING_SNAKE_CASE_						 = 1
									SCREAMING_SNAKE_CASE_						 = 15

									# dummy input_ids and scores
									SCREAMING_SNAKE_CASE_						 = ids_tensor((batch_size, sequence_length)       ,       lowercase__    )
									SCREAMING_SNAKE_CASE_						 = input_ids.copy()

									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(lowercase__       ,       lowercase__    )
									SCREAMING_SNAKE_CASE_						 = scores.copy()

									# instantiate all dist processors
									SCREAMING_SNAKE_CASE_						 = FlaxTemperatureLogitsWarper(temperature=0.5    )
									SCREAMING_SNAKE_CASE_						 = FlaxTopKLogitsWarper(3    )
									SCREAMING_SNAKE_CASE_						 = FlaxTopPLogitsWarper(0.8    )

									# instantiate all logits processors
									SCREAMING_SNAKE_CASE_						 = FlaxMinLengthLogitsProcessor(min_length=10       ,       eos_token_id=lowercase__    )
									SCREAMING_SNAKE_CASE_						 = FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowercase__    )
									SCREAMING_SNAKE_CASE_						 = FlaxForcedEOSTokenLogitsProcessor(max_length=lowercase__       ,       eos_token_id=lowercase__    )

									SCREAMING_SNAKE_CASE_						 = 10

									# no processor list
									SCREAMING_SNAKE_CASE_						 = temp_dist_warp(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
									SCREAMING_SNAKE_CASE_						 = top_k_warp(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
									SCREAMING_SNAKE_CASE_						 = top_p_warp(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
									SCREAMING_SNAKE_CASE_						 = min_dist_proc(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
									SCREAMING_SNAKE_CASE_						 = bos_dist_proc(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
									SCREAMING_SNAKE_CASE_						 = eos_dist_proc(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )

									# with processor list
									SCREAMING_SNAKE_CASE_						 = FlaxLogitsProcessorList(
									    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]    )
									SCREAMING_SNAKE_CASE_						 = processor(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )

									# scores should be equal
									self.assertTrue(jnp.allclose(lowercase__       ,       lowercase__       ,       atol=1e-3    )    )

									# input_ids should never be changed
									self.assertListEqual(input_ids.tolist()       ,       input_ids_comp.tolist()    )







							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = 4
									SCREAMING_SNAKE_CASE_						 = 10
									SCREAMING_SNAKE_CASE_						 = 15
									SCREAMING_SNAKE_CASE_						 = 2
									SCREAMING_SNAKE_CASE_						 = 1
									SCREAMING_SNAKE_CASE_						 = 15

									# dummy input_ids and scores
									SCREAMING_SNAKE_CASE_						 = ids_tensor((batch_size, sequence_length)       ,       lowercase__    )
									SCREAMING_SNAKE_CASE_						 = input_ids.copy()

									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(lowercase__       ,       lowercase__    )
									SCREAMING_SNAKE_CASE_						 = scores.copy()

									# instantiate all dist processors
									SCREAMING_SNAKE_CASE_						 = FlaxTemperatureLogitsWarper(temperature=0.5    )
									SCREAMING_SNAKE_CASE_						 = FlaxTopKLogitsWarper(3    )
									SCREAMING_SNAKE_CASE_						 = FlaxTopPLogitsWarper(0.8    )

									# instantiate all logits processors
									SCREAMING_SNAKE_CASE_						 = FlaxMinLengthLogitsProcessor(min_length=10       ,       eos_token_id=lowercase__    )
									SCREAMING_SNAKE_CASE_						 = FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowercase__    )
									SCREAMING_SNAKE_CASE_						 = FlaxForcedEOSTokenLogitsProcessor(max_length=lowercase__       ,       eos_token_id=lowercase__    )

									SCREAMING_SNAKE_CASE_						 = 10

									# no processor list
									def run_no_processor_list(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):
											SCREAMING_SNAKE_CASE_						 = temp_dist_warp(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
											SCREAMING_SNAKE_CASE_						 = top_k_warp(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
											SCREAMING_SNAKE_CASE_						 = top_p_warp(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
											SCREAMING_SNAKE_CASE_						 = min_dist_proc(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
											SCREAMING_SNAKE_CASE_						 = bos_dist_proc(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
											SCREAMING_SNAKE_CASE_						 = eos_dist_proc(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
											return scores

									# with processor list
									def run_processor_list(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):
											SCREAMING_SNAKE_CASE_						 = FlaxLogitsProcessorList(
											    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]    )
											SCREAMING_SNAKE_CASE_						 = processor(lowercase__       ,       lowercase__       ,       cur_len=lowercase__    )
											return scores

									SCREAMING_SNAKE_CASE_						 = jax.jit(lowercase__    )
									SCREAMING_SNAKE_CASE_						 = jax.jit(lowercase__    )

									SCREAMING_SNAKE_CASE_						 = jitted_run_no_processor_list(lowercase__       ,       lowercase__       ,       lowercase__    )
									SCREAMING_SNAKE_CASE_						 = jitted_run_processor_list(lowercase__       ,       lowercase__       ,       lowercase__    )

									# scores should be equal
									self.assertTrue(jnp.allclose(lowercase__       ,       lowercase__       ,       atol=1e-3    )    )

									# input_ids should never be changed
									self.assertListEqual(input_ids.tolist()       ,       input_ids_comp.tolist()    )

"""simple docstring"""




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 ( __a    ):  # picklable for multiprocessing
		return x.sum()



def        _lowerCamelCase ( __a    ):  # picklable for multiprocessing
		return i + 1





@dataclass
class  snake_case     :
							UpperCAmelCase__						     =      42
							UpperCAmelCase__						     =      42







class  snake_case     (							__lowercase   ):

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = {}
									SCREAMING_SNAKE_CASE_						 = []
									SCREAMING_SNAKE_CASE_						 = 1
									SCREAMING_SNAKE_CASE_						 = [1, 2]
									SCREAMING_SNAKE_CASE_						 = {'''a''': 1, '''b''': 2}
									SCREAMING_SNAKE_CASE_						 = {'''a''': [1, 2], '''b''': [3, 4]}
									SCREAMING_SNAKE_CASE_						 = {'''a''': {'''1''': 1}, '''b''': 2}
									SCREAMING_SNAKE_CASE_						 = {'''a''': 1, '''b''': 2, '''c''': 3, '''d''': 4}
									SCREAMING_SNAKE_CASE_						 = {}
									SCREAMING_SNAKE_CASE_						 = []
									SCREAMING_SNAKE_CASE_						 = 2
									SCREAMING_SNAKE_CASE_						 = [2, 3]
									SCREAMING_SNAKE_CASE_						 = {'''a''': 2, '''b''': 3}
									SCREAMING_SNAKE_CASE_						 = {'''a''': [2, 3], '''b''': [4, 5]}
									SCREAMING_SNAKE_CASE_						 = {'''a''': {'''1''': 2}, '''b''': 3}
									SCREAMING_SNAKE_CASE_						 = {'''a''': 2, '''b''': 3, '''c''': 4, '''d''': 5}
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )

									SCREAMING_SNAKE_CASE_						 = 2
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )

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

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = {'''a''': 1, '''b''': 2}
									SCREAMING_SNAKE_CASE_						 = {'''a''': 3, '''b''': 4}
									SCREAMING_SNAKE_CASE_						 = {'''a''': 5, '''b''': 6}
									SCREAMING_SNAKE_CASE_						 = sorted([('''a''', (1, 3, 5)), ('''b''', (2, 4, 6))]    )
									self.assertEqual(sorted(zip_dict(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )    )       ,       SCREAMING_SNAKE_CASE_    )







							def _lowercase					(self    ):

									"""simple docstring"""







									class  snake_case     :
																UpperCAmelCase__						     =      '''bar'''







									SCREAMING_SNAKE_CASE_						 = Foo()
									self.assertEqual(foo.my_attr       ,       '''bar'''    )
									with temporary_assignment(SCREAMING_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 ( __a,							__a,							__a    ):
		with patch('''datasets.utils.py_utils._single_map_nested'''    ) as mock_single_map_nested, patch(
		    '''datasets.parallel.parallel.Pool'''    ) as mock_multiprocessing_pool:
				SCREAMING_SNAKE_CASE_						 = {F'{i}': i for i in range(__a    )}
				SCREAMING_SNAKE_CASE_						 = map_nested(lambda __a    : x + 10,							__a,							num_proc=__a,							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  snake_case     (							__lowercase   ):

							@require_tf
							def _lowercase					(self    ):

									"""simple docstring"""







									import tensorflow as tf
									from tensorflow.keras import layers

									SCREAMING_SNAKE_CASE_						 = layers.Dense(2    )

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

									with temp_seed(42       ,       set_tensorflow=SCREAMING_SNAKE_CASE_    ):
											SCREAMING_SNAKE_CASE_						 = gen_random_output()
									with temp_seed(42       ,       set_tensorflow=SCREAMING_SNAKE_CASE_    ):
											SCREAMING_SNAKE_CASE_						 = gen_random_output()
									SCREAMING_SNAKE_CASE_						 = gen_random_output()

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

							@require_torch
							def _lowercase					(self    ):

									"""simple docstring"""







									import torch

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

									with temp_seed(42       ,       set_pytorch=SCREAMING_SNAKE_CASE_    ):
											SCREAMING_SNAKE_CASE_						 = gen_random_output()
									with temp_seed(42       ,       set_pytorch=SCREAMING_SNAKE_CASE_    ):
											SCREAMING_SNAKE_CASE_						 = gen_random_output()
									SCREAMING_SNAKE_CASE_						 = gen_random_output()

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







							def _lowercase					(self    ):

									"""simple docstring"""







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

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

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







@pytest.mark.parametrize('''input_data''',							[{}]    )
def        _lowerCamelCase ( __a    ):
		SCREAMING_SNAKE_CASE_						 = NestedDataStructure(__a    ).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 ( __a,							__a    ):
		SCREAMING_SNAKE_CASE_						 = NestedDataStructure(__a    ).flatten()
		assert output == expected_output



def        _lowerCamelCase ( ):
		SCREAMING_SNAKE_CASE_						 = A(x=1,							y='''foobar'''    )
		SCREAMING_SNAKE_CASE_						 = {'''x''': 1, '''y''': '''foobar'''}
		assert asdict(__a    ) == expected_output

		SCREAMING_SNAKE_CASE_						 = {'''a''': {'''b''': A(x=10,							y='''foo'''    )}, '''c''': [A(x=20,							y='''bar'''    )]}
		SCREAMING_SNAKE_CASE_						 = {'''a''': {'''b''': {'''x''': 10, '''y''': '''foo'''}}, '''c''': [{'''x''': 20, '''y''': '''bar'''}]}
		assert asdict(__a    ) == expected_output

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



def        _lowerCamelCase ( __a    ):
		return text.split()



def        _lowerCamelCase ( __a    ):
		yield (time.time(), content)
		time.sleep(2    )
		yield (time.time(), content)



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

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

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