





'''simple docstring'''






import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotConfig, is_flax_available
from transformers.testing_utils import jax_device, 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
						__lowerCamelCase						: Dict		=   'platform'

						import jax
						import jax.numpy as jnp

						from transformers import BlenderbotTokenizer
						from transformers.models.blenderbot.modeling_flax_blenderbot import (
						    FlaxBlenderbotForConditionalGeneration,
						    FlaxBlenderbotModel,
						    shift_tokens_right,
						)


def 					_a   (__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       			=np.where(input_ids != config.pad_token_id  ,      1  ,      0   )
		if decoder_attention_mask is None:
				_UpperCamelCase       			=np.where(decoder_input_ids != config.pad_token_id  ,      1  ,      0   )
		if head_mask is None:
				_UpperCamelCase       			=np.ones((config.encoder_layers, config.encoder_attention_heads)   )
		if decoder_head_mask is None:
				_UpperCamelCase       			=np.ones((config.decoder_layers, config.decoder_attention_heads)   )
		if cross_attn_head_mask is None:
				_UpperCamelCase       			=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 		UpperCAmelCase   :


		"""simple docstring"""







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


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

				_UpperCamelCase       			=shift_tokens_right(UpperCamelCase__							,	1							,	2  )

				_UpperCamelCase       			=BlenderbotConfig(
				    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=UpperCamelCase__							,	)
				_UpperCamelCase       			=prepare_blenderbot_inputs_dict(UpperCamelCase__							,	UpperCamelCase__							,	UpperCamelCase__  )
				return config, inputs_dict


		def   UpperCamelCase__      (				self    :					str  )     ->						Any:
				_UpperCamelCase			,     _UpperCamelCase       			=self.prepare_config_and_inputs()
				return config, inputs_dict


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

				_UpperCamelCase       			=model.encode(inputs_dict['''input_ids''']  )

				_UpperCamelCase			,     _UpperCamelCase       			=(
				    inputs_dict['''decoder_input_ids'''],
				    inputs_dict['''decoder_attention_mask'''],
				)

				_UpperCamelCase       			=model.init_cache(decoder_input_ids.shape[0]							,	UpperCamelCase__							,	UpperCamelCase__  )
				_UpperCamelCase       			=jnp.ones((decoder_input_ids.shape[0], max_decoder_length)							,	dtype='''i4'''  )

				_UpperCamelCase       			=jnp.broadcast_to(
				    jnp.arange(decoder_input_ids.shape[-1] - 1  )[None, :]							,	(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)							,	)
				_UpperCamelCase       			=model.decode(
				    decoder_input_ids[:, :-1]							,	UpperCamelCase__							,	decoder_attention_mask=UpperCamelCase__							,	past_key_values=UpperCamelCase__							,	decoder_position_ids=UpperCamelCase__							,	)

				_UpperCamelCase       			=jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]							,	dtype='''i4'''  )
				_UpperCamelCase       			=model.decode(
				    decoder_input_ids[:, -1:]							,	UpperCamelCase__							,	decoder_attention_mask=UpperCamelCase__							,	past_key_values=outputs_cache.past_key_values							,	decoder_position_ids=UpperCamelCase__							,	)

				_UpperCamelCase       			=model.decode(UpperCamelCase__							,	UpperCamelCase__  )

				_UpperCamelCase       			=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    :					str							,	UpperCamelCase__    :					Any							,	UpperCamelCase__    :					Any							,	UpperCamelCase__    :					Dict  )     ->						Any:
				_UpperCamelCase       			=20
				_UpperCamelCase       			=model_class_name(UpperCamelCase__  )

				_UpperCamelCase       			=model.encode(inputs_dict['''input_ids''']  )

				_UpperCamelCase			,     _UpperCamelCase       			=(
				    inputs_dict['''decoder_input_ids'''],
				    inputs_dict['''decoder_attention_mask'''],
				)

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

				_UpperCamelCase       			=model.init_cache(decoder_input_ids.shape[0]							,	UpperCamelCase__							,	UpperCamelCase__  )
				_UpperCamelCase       			=jnp.broadcast_to(
				    jnp.arange(decoder_input_ids.shape[-1] - 1  )[None, :]							,	(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)							,	)

				_UpperCamelCase       			=model.decode(
				    decoder_input_ids[:, :-1]							,	UpperCamelCase__							,	decoder_attention_mask=UpperCamelCase__							,	past_key_values=UpperCamelCase__							,	decoder_position_ids=UpperCamelCase__							,	)
				_UpperCamelCase       			=jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]							,	dtype='''i4'''  )
				_UpperCamelCase       			=model.decode(
				    decoder_input_ids[:, -1:]							,	UpperCamelCase__							,	past_key_values=outputs_cache.past_key_values							,	decoder_attention_mask=UpperCamelCase__							,	decoder_position_ids=UpperCamelCase__							,	)

				_UpperCamelCase       			=model.decode(UpperCamelCase__							,	UpperCamelCase__							,	decoder_attention_mask=UpperCamelCase__  )

				_UpperCamelCase       			=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 		UpperCAmelCase   (				unittest.TestCase):


		"""simple docstring"""







		lowerCAmelCase_					=				99


		def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						Tuple:
				_UpperCamelCase       			=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       			=input_ids.shape[0]
				_UpperCamelCase       			=BlenderbotConfig(
				    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    :					Tuple  )     ->						List[str]:
				_UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase       			=self._get_config_and_data()
				_UpperCamelCase       			=FlaxBlenderbotForConditionalGeneration(UpperCamelCase__  )
				_UpperCamelCase       			=lm_model(input_ids=UpperCamelCase__  )
				_UpperCamelCase       			=(batch_size, input_ids.shape[1], config.vocab_size)
				self.assertEqual(outputs['''logits'''].shape							,	UpperCamelCase__  )


		def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						List[Any]:
				_UpperCamelCase       			=BlenderbotConfig(
				    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       			=FlaxBlenderbotForConditionalGeneration(UpperCamelCase__  )
				_UpperCamelCase       			=np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]							,	dtype=np.intaa  )
				_UpperCamelCase       			=np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]							,	dtype=np.intaa  )
				_UpperCamelCase       			=lm_model(input_ids=UpperCamelCase__							,	decoder_input_ids=UpperCamelCase__  )
				_UpperCamelCase       			=(*summary.shape, config.vocab_size)
				self.assertEqual(outputs['''logits'''].shape							,	UpperCamelCase__  )




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







@require_flax
class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase   ,     lowercase_):


		"""simple docstring"""







		lowerCAmelCase_					=				True
		lowerCAmelCase_					=				(
		    (
		        FlaxBlenderbotModel,
		        FlaxBlenderbotForConditionalGeneration,
		    )
		    if is_flax_available()
		    else ()
		)
		lowerCAmelCase_					=				(FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()


		def   UpperCamelCase__      (				self    :					Tuple  )     ->						Any:
				_UpperCamelCase       			=FlaxBlenderbotModelTester(self  )


		def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						Any:
				_UpperCamelCase			,     _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
				for model_class in self.all_model_classes:
						self.model_tester.check_use_cache_forward(UpperCamelCase__							,	UpperCamelCase__							,	UpperCamelCase__  )


		def   UpperCamelCase__      (				self    :					int  )     ->						Any:
				_UpperCamelCase			,     _UpperCamelCase       			=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(UpperCamelCase__							,	UpperCamelCase__							,	UpperCamelCase__  )


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

				for model_class in self.all_model_classes:
						with self.subTest(model_class.__name__  ):
								_UpperCamelCase       			=self._prepare_for_class(UpperCamelCase__							,	UpperCamelCase__  )
								_UpperCamelCase       			=model_class(UpperCamelCase__  )

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

								with self.subTest('''JIT Enabled'''  ):
										_UpperCamelCase       			=encode_jitted(**UpperCamelCase__  ).to_tuple()

								with self.subTest('''JIT Disabled'''  ):
										with jax.disable_jit():
												_UpperCamelCase       			=encode_jitted(**UpperCamelCase__  ).to_tuple()

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


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

				for model_class in self.all_model_classes:
						with self.subTest(model_class.__name__  ):
								_UpperCamelCase       			=model_class(UpperCamelCase__  )
								_UpperCamelCase       			=model.encode(inputs_dict['''input_ids''']							,	inputs_dict['''attention_mask''']  )

								_UpperCamelCase       			={
								    '''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(UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					Tuple  ):
										return model.decode(
										    decoder_input_ids=UpperCamelCase__							,	decoder_attention_mask=UpperCamelCase__							,	encoder_outputs=UpperCamelCase__							,	)

								with self.subTest('''JIT Enabled'''  ):
										_UpperCamelCase       			=decode_jitted(**UpperCamelCase__  ).to_tuple()

								with self.subTest('''JIT Disabled'''  ):
										with jax.disable_jit():
												_UpperCamelCase       			=decode_jitted(**UpperCamelCase__  ).to_tuple()

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


		@slow
		def   UpperCamelCase__      (				self    :					List[str]  )     ->						Dict:
				for model_class_name in self.all_model_classes:
						_UpperCamelCase       			=model_class_name.from_pretrained('''facebook/blenderbot-400M-distill'''  )
						# FlaxBlenderbotForSequenceClassification expects eos token in input_ids
						_UpperCamelCase       			=np.ones((1, 1)  ) * model.config.eos_token_id
						_UpperCamelCase       			=model(UpperCamelCase__  )
						self.assertIsNotNone(UpperCamelCase__  )




		@unittest.skipUnless(jax_device != '''cpu'''							,	'''3B test too slow on CPU.'''  )
		@slow
		def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						List[str]:
				_UpperCamelCase       			={'''num_beams''': 1, '''early_stopping''': True, '''min_length''': 15, '''max_length''': 25}
				_UpperCamelCase       			={'''skip_special_tokens''': True, '''clean_up_tokenization_spaces''': True}

				_UpperCamelCase       			=FlaxBlenderbotForConditionalGeneration.from_pretrained('''facebook/blenderbot-3B'''							,	from_pt=UpperCamelCase__  )
				_UpperCamelCase       			=BlenderbotTokenizer.from_pretrained('''facebook/blenderbot-3B'''  )

				_UpperCamelCase       			=['''Sam''']
				_UpperCamelCase       			=tokenizer(UpperCamelCase__							,	return_tensors='''jax'''  )

				_UpperCamelCase       			=model.generate(**UpperCamelCase__							,	**UpperCamelCase__  )
				_UpperCamelCase       			='''Sam is a great name. It means "sun" in Gaelic.'''

				_UpperCamelCase       			=tokenizer.batch_decode(UpperCamelCase__							,	**UpperCamelCase__  )
				assert generated_txt[0].strip() == tgt_text













'''simple docstring'''






import contextlib
from multiprocessing import Pool, RLock

from tqdm.auto import tqdm

from ..utils import experimental, logging


__lowerCamelCase						: Union[str, Any]		=   logging.get_logger(__name__)







class 		UpperCAmelCase   :


  """simple docstring"""







  lowerCAmelCase_					=				None




@experimental
def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  if ParallelBackendConfig.backend_name is None:
    return _map_with_multiprocessing_pool(
        __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )

  return _map_with_joblib(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  _UpperCamelCase       			=num_proc if num_proc <= len(__SCREAMING_SNAKE_CASE   ) else len(__SCREAMING_SNAKE_CASE   )
  _UpperCamelCase       			=[]  # We organize the splits ourselve (contiguous splits)
  for index in range(__SCREAMING_SNAKE_CASE   ):
    _UpperCamelCase       			=len(__SCREAMING_SNAKE_CASE   ) // num_proc
    _UpperCamelCase       			=len(__SCREAMING_SNAKE_CASE   ) % num_proc
    _UpperCamelCase       			=div * index + min(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )
    _UpperCamelCase       			=start + div + (1 if index < mod else 0)
    split_kwds.append((function, iterable[start:end], types, index, disable_tqdm, desc)   )

  if len(__SCREAMING_SNAKE_CASE   ) != sum(len(i[1]   ) for i in split_kwds   ):
    raise ValueError(
        f'''Error dividing inputs iterable among processes. '''
        f'''Total number of objects {len(__SCREAMING_SNAKE_CASE   )}, '''
        f'''length: {sum(len(i[1]   ) for i in split_kwds   )}'''   )

  logger.info(
      f'''Spawning {num_proc} processes for {len(__SCREAMING_SNAKE_CASE   )} objects in slices of {[len(i[1]   ) for i in split_kwds]}'''   )
  _UpperCamelCase			,     _UpperCamelCase       			=None, None
  if not disable_tqdm:
    _UpperCamelCase			,     _UpperCamelCase       			=(RLock(),), tqdm.set_lock
  with Pool(__SCREAMING_SNAKE_CASE  ,      initargs=__SCREAMING_SNAKE_CASE  ,      initializer=__SCREAMING_SNAKE_CASE   ) as pool:
    _UpperCamelCase       			=pool.map(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )
  logger.info(f'''Finished {num_proc} processes'''   )
  _UpperCamelCase       			=[obj for proc_res in mapped for obj in proc_res]
  logger.info(f'''Unpacked {len(__SCREAMING_SNAKE_CASE   )} objects'''   )

  return mapped


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  import joblib

  with joblib.parallel_backend(ParallelBackendConfig.backend_name  ,      n_jobs=__SCREAMING_SNAKE_CASE   ):
    return joblib.Parallel()(
        joblib.delayed(__SCREAMING_SNAKE_CASE   )((function, obj, types, None, True, None)   ) for obj in iterable   )


@experimental
@contextlib.contextmanager
def 					_a   (__SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  _UpperCamelCase       			=backend_name

  if backend_name == "spark":
    from joblibspark import register_spark

    register_spark()

    # TODO: call create_cache_and_write_probe if "download" in steps
    # TODO: raise NotImplementedError when Dataset.map etc is called

  try:
    yield
  finally:
    _UpperCamelCase       			=None






