






"""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
  UpperCAmelCase_     :				Tuple         =				"""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    (__a					,       __a					,       __a=None					,       __a=None					,       __a=None					,       __a=None					,       __a=None					,       __a=None					,       )				-> Union[str, Any]:
    """simple docstring"""






    if attention_mask is None:
        SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	np.where(input_ids != config.pad_token_id					,       1					,       0							)
    if decoder_attention_mask is None:
        SCREAMING_SNAKE_CASE_   :						Tuple							=	np.where(decoder_input_ids != config.pad_token_id					,       1					,       0							)
    if head_mask is None:
        SCREAMING_SNAKE_CASE_   :						Dict							=	np.ones((config.encoder_layers, config.encoder_attention_heads)							)
    if decoder_head_mask is None:
        SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	np.ones((config.decoder_layers, config.decoder_attention_heads)							)
    if cross_attn_head_mask is None:
        SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	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  :						Union[str, Any]		,			lowercase_  :						int		,			lowercase_  :						Optional[int]=13		,			lowercase_  :						str=7		,			lowercase_  :						Tuple=True		,			lowercase_  :						Optional[int]=False		,			lowercase_  :						Tuple=99		,			lowercase_  :						List[str]=16		,			lowercase_  :						List[Any]=2		,			lowercase_  :						Union[str, Any]=4		,			lowercase_  :						List[str]=4		,			lowercase_  :						List[Any]="gelu"		,			lowercase_  :						List[str]=0.1		,			lowercase_  :						int=0.1		,			lowercase_  :						List[str]=32		,			lowercase_  :						List[str]=2		,			lowercase_  :						int=1		,			lowercase_  :						Dict=0		,			lowercase_  :						Optional[int]=0.02		,			):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Any							=	parent
          SCREAMING_SNAKE_CASE_   :						int							=	batch_size
          SCREAMING_SNAKE_CASE_   :						Dict							=	seq_length
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	is_training
          SCREAMING_SNAKE_CASE_   :						int							=	use_labels
          SCREAMING_SNAKE_CASE_   :						List[str]							=	vocab_size
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	hidden_size
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	num_hidden_layers
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	num_attention_heads
          SCREAMING_SNAKE_CASE_   :						Any							=	intermediate_size
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	hidden_act
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	hidden_dropout_prob
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	attention_probs_dropout_prob
          SCREAMING_SNAKE_CASE_   :						Any							=	max_position_embeddings
          SCREAMING_SNAKE_CASE_   :						Any							=	eos_token_id
          SCREAMING_SNAKE_CASE_   :						List[str]							=	pad_token_id
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	bos_token_id
          SCREAMING_SNAKE_CASE_   :						Tuple							=	initializer_range
      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						List[Any]							=	np.clip(ids_tensor([self.batch_size, self.seq_length - 1]		,			self.vocab_size)		,			3		,			self.vocab_size)
          SCREAMING_SNAKE_CASE_   :						Dict							=	np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)		,			dtype=np.intaa))		,			-1)

          SCREAMING_SNAKE_CASE_   :						str							=	shift_tokens_right(lowercase_		,			1		,			2)

          SCREAMING_SNAKE_CASE_   :						List[Any]							=	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=lowercase_		,			)
          SCREAMING_SNAKE_CASE_   :						Dict							=	prepare_blenderbot_inputs_dict(lowercase_		,			lowercase_		,			lowercase_)
          return config, inputs_dict
      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						int							=	self.prepare_config_and_inputs()
          return config, inputs_dict
      def _SCREAMING_SNAKE_CASE				(		self  :						Union[str, Any]		,			lowercase_  :						Union[str, Any]		,			lowercase_  :						List[str]		,			lowercase_  :						List[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						int							=	20
          SCREAMING_SNAKE_CASE_   :						Dict							=	model_class_name(lowercase_)

          SCREAMING_SNAKE_CASE_   :						int							=	model.encode(inputs_dict['''input_ids'''])

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

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

          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	jnp.broadcast_to(
              jnp.arange(decoder_input_ids.shape[-1] - 1)[None, :]		,			(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)		,			)
          SCREAMING_SNAKE_CASE_   :						List[str]							=	model.decode(
              decoder_input_ids[:, :-1]		,			lowercase_		,			decoder_attention_mask=lowercase_		,			past_key_values=lowercase_		,			decoder_position_ids=lowercase_		,			)

          SCREAMING_SNAKE_CASE_   :						Tuple							=	jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]		,			dtype='''i4''')
          SCREAMING_SNAKE_CASE_   :						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_		,			)

          SCREAMING_SNAKE_CASE_   :						str							=	model.decode(lowercase_		,			lowercase_)

          SCREAMING_SNAKE_CASE_   :						Optional[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}')



      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[Any]		,			lowercase_  :						Optional[Any]		,			lowercase_  :						List[str]		,			lowercase_  :						List[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	20
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	model_class_name(lowercase_)

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

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

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

          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	model.init_cache(decoder_input_ids.shape[0]		,			lowercase_		,			lowercase_)
          SCREAMING_SNAKE_CASE_   :						Tuple							=	jnp.broadcast_to(
              jnp.arange(decoder_input_ids.shape[-1] - 1)[None, :]		,			(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)		,			)

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	model.decode(
              decoder_input_ids[:, :-1]		,			lowercase_		,			decoder_attention_mask=lowercase_		,			past_key_values=lowercase_		,			decoder_position_ids=lowercase_		,			)
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]		,			dtype='''i4''')
          SCREAMING_SNAKE_CASE_   :						Dict							=	model.decode(
              decoder_input_ids[:, -1:]		,			lowercase_		,			past_key_values=outputs_cache.past_key_values		,			decoder_attention_mask=lowercase_		,			decoder_position_ids=lowercase_		,			)

          SCREAMING_SNAKE_CASE_   :						str							=	model.decode(lowercase_		,			lowercase_		,			decoder_attention_mask=lowercase_)

          SCREAMING_SNAKE_CASE_   :						str							=	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'''

      __UpperCamelCase							=						9_9
      def _SCREAMING_SNAKE_CASE				(		self  :						int):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						int							=	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		,			)

          SCREAMING_SNAKE_CASE_   :						int							=	input_ids.shape[0]
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	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 _SCREAMING_SNAKE_CASE				(		self  :						List[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						List[Any]							=	self._get_config_and_data()
          SCREAMING_SNAKE_CASE_   :						str							=	FlaxBlenderbotForConditionalGeneration(lowercase_)
          SCREAMING_SNAKE_CASE_   :						str							=	lm_model(input_ids=lowercase_)
          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	(batch_size, input_ids.shape[1], config.vocab_size)
          self.assertEqual(outputs['''logits'''].shape		,			lowercase_)
      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]):




          '''simple docstring'''


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



      def _SCREAMING_SNAKE_CASE				(		self  :						List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						List[Any]							=	np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]		,			dtype=np.intaa)
          SCREAMING_SNAKE_CASE_   :						int							=	shift_tokens_right(lowercase_		,			1		,			2)
          SCREAMING_SNAKE_CASE_   :						Tuple							=	np.equal(lowercase_		,			1).astype(np.floataa).sum()
          SCREAMING_SNAKE_CASE_   :						List[str]							=	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'''

      __UpperCamelCase							=						True
      __UpperCamelCase							=						(
          (
              FlaxBlenderbotModel,
              FlaxBlenderbotForConditionalGeneration,
          )
          if is_flax_available()
          else ()
      )
      __UpperCamelCase							=						(FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	FlaxBlenderbotModelTester(self)
      def _SCREAMING_SNAKE_CASE				(		self  :						Union[str, Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						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 _SCREAMING_SNAKE_CASE				(		self  :						List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						Dict							=	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 _SCREAMING_SNAKE_CASE				(		self  :						Dict):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
              with self.subTest(model_class.__name__):
                  SCREAMING_SNAKE_CASE_   :						Tuple							=	self._prepare_for_class(lowercase_		,			lowercase_)
                  SCREAMING_SNAKE_CASE_   :						Any							=	model_class(lowercase_)

                  @jax.jit
                  def encode_jitted(lowercase_  :						List[Any]		,			lowercase_  :						Any=None		,			**lowercase_  :						Optional[Any]):
                      return model.encode(input_ids=lowercase_		,			attention_mask=lowercase_)

                  with self.subTest('''JIT Enabled'''):
                      SCREAMING_SNAKE_CASE_   :						Any							=	encode_jitted(**lowercase_).to_tuple()

                  with self.subTest('''JIT Disabled'''):
                      with jax.disable_jit():
                          SCREAMING_SNAKE_CASE_   :						int							=	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 _SCREAMING_SNAKE_CASE				(		self  :						Dict):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	self.model_tester.prepare_config_and_inputs_for_common()

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

                  SCREAMING_SNAKE_CASE_   :						List[str]							=	{
                      '''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_  :						int		,			lowercase_  :						Dict		,			lowercase_  :						Optional[int]):
                      return model.decode(
                          decoder_input_ids=lowercase_		,			decoder_attention_mask=lowercase_		,			encoder_outputs=lowercase_		,			)

                  with self.subTest('''JIT Enabled'''):
                      SCREAMING_SNAKE_CASE_   :						Dict							=	decode_jitted(**lowercase_).to_tuple()

                  with self.subTest('''JIT Disabled'''):
                      with jax.disable_jit():
                          SCREAMING_SNAKE_CASE_   :						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 _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]):




          '''simple docstring'''


          for model_class_name in self.all_model_classes:
              SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	model_class_name.from_pretrained('''facebook/blenderbot-400M-distill''')
              # FlaxBlenderbotForSequenceClassification expects eos token in input_ids
              SCREAMING_SNAKE_CASE_   :						List[str]							=	np.ones((1, 1)) * model.config.eos_token_id
              SCREAMING_SNAKE_CASE_   :						str							=	model(lowercase_)
              self.assertIsNotNone(lowercase_)



      @unittest.skipUnless(jax_device != '''cpu'''		,			'''3B test too slow on CPU.''')
      @slow
      def _SCREAMING_SNAKE_CASE				(		self  :						int):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						int							=	{'''num_beams''': 1, '''early_stopping''': True, '''min_length''': 15, '''max_length''': 25}
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	{'''skip_special_tokens''': True, '''clean_up_tokenization_spaces''': True}

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	FlaxBlenderbotForConditionalGeneration.from_pretrained('''facebook/blenderbot-3B'''		,			from_pt=lowercase_)
          SCREAMING_SNAKE_CASE_   :						Any							=	BlenderbotTokenizer.from_pretrained('''facebook/blenderbot-3B''')

          SCREAMING_SNAKE_CASE_   :						List[str]							=	['''Sam''']
          SCREAMING_SNAKE_CASE_   :						str							=	tokenizer(lowercase_		,			return_tensors='''jax''')

          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	model.generate(**lowercase_		,			**lowercase_)
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	'''Sam is a great name. It means "sun" in Gaelic.'''

          SCREAMING_SNAKE_CASE_   :						int							=	tokenizer.batch_decode(lowercase_		,			**lowercase_)
          assert generated_txt[0].strip() == tgt_text










"""simple docstring"""






import unittest

import numpy as np

from transformers.testing_utils import require_flax, require_tf, require_torch
from transformers.utils import (
    expand_dims,
    flatten_dict,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    reshape,
    squeeze,
    transpose,
)


if is_flax_available():
  import jax.numpy as jnp

if is_tf_available():
  import tensorflow as tf

if is_torch_available():
  import torch




class 							lowerCAmelCase__       (    unittest.TestCase				):






      '''simple docstring'''

      def _SCREAMING_SNAKE_CASE				(		self  :						str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						str							=	{
              '''task_specific_params''': {
                  '''summarization''': {'''length_penalty''': 1.0, '''max_length''': 128, '''min_length''': 12, '''num_beams''': 4},
                  '''summarization_cnn''': {'''length_penalty''': 2.0, '''max_length''': 142, '''min_length''': 56, '''num_beams''': 4},
                  '''summarization_xsum''': {'''length_penalty''': 1.0, '''max_length''': 62, '''min_length''': 11, '''num_beams''': 6},
              }
          }
          SCREAMING_SNAKE_CASE_   :						List[str]							=	{
              '''task_specific_params.summarization.length_penalty''': 1.0,
              '''task_specific_params.summarization.max_length''': 128,
              '''task_specific_params.summarization.min_length''': 12,
              '''task_specific_params.summarization.num_beams''': 4,
              '''task_specific_params.summarization_cnn.length_penalty''': 2.0,
              '''task_specific_params.summarization_cnn.max_length''': 142,
              '''task_specific_params.summarization_cnn.min_length''': 56,
              '''task_specific_params.summarization_cnn.num_beams''': 4,
              '''task_specific_params.summarization_xsum.length_penalty''': 1.0,
              '''task_specific_params.summarization_xsum.max_length''': 62,
              '''task_specific_params.summarization_xsum.min_length''': 11,
              '''task_specific_params.summarization_xsum.num_beams''': 6,
          }

          self.assertEqual(flatten_dict(lowercase_)		,			lowercase_)
      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	np.random.randn(3		,			4)
          self.assertTrue(np.allclose(transpose(lowercase_)		,			x.transpose()))

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	np.random.randn(3		,			4		,			5)
          self.assertTrue(np.allclose(transpose(lowercase_		,			axes=(1, 2, 0))		,			x.transpose((1, 2, 0))))
      @require_torch
      def _SCREAMING_SNAKE_CASE				(		self  :						Dict):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						str							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						str							=	torch.tensor(lowercase_)
          self.assertTrue(np.allclose(transpose(lowercase_)		,			transpose(lowercase_).numpy()))

          SCREAMING_SNAKE_CASE_   :						List[Any]							=	np.random.randn(3		,			4		,			5)
          SCREAMING_SNAKE_CASE_   :						Dict							=	torch.tensor(lowercase_)
          self.assertTrue(np.allclose(transpose(lowercase_		,			axes=(1, 2, 0))		,			transpose(lowercase_		,			axes=(1, 2, 0)).numpy()))
      @require_tf
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						str							=	tf.constant(lowercase_)
          self.assertTrue(np.allclose(transpose(lowercase_)		,			transpose(lowercase_).numpy()))

          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	np.random.randn(3		,			4		,			5)
          SCREAMING_SNAKE_CASE_   :						int							=	tf.constant(lowercase_)
          self.assertTrue(np.allclose(transpose(lowercase_		,			axes=(1, 2, 0))		,			transpose(lowercase_		,			axes=(1, 2, 0)).numpy()))
      @require_flax
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Tuple							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						Any							=	jnp.array(lowercase_)
          self.assertTrue(np.allclose(transpose(lowercase_)		,			np.asarray(transpose(lowercase_))))

          SCREAMING_SNAKE_CASE_   :						Dict							=	np.random.randn(3		,			4		,			5)
          SCREAMING_SNAKE_CASE_   :						Any							=	jnp.array(lowercase_)
          self.assertTrue(np.allclose(transpose(lowercase_		,			axes=(1, 2, 0))		,			np.asarray(transpose(lowercase_		,			axes=(1, 2, 0)))))
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Dict							=	np.random.randn(3		,			4)
          self.assertTrue(np.allclose(reshape(lowercase_		,			(4, 3))		,			np.reshape(lowercase_		,			(4, 3))))

          SCREAMING_SNAKE_CASE_   :						Dict							=	np.random.randn(3		,			4		,			5)
          self.assertTrue(np.allclose(reshape(lowercase_		,			(12, 5))		,			np.reshape(lowercase_		,			(12, 5))))
      @require_torch
      def _SCREAMING_SNAKE_CASE				(		self  :						Union[str, Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						List[str]							=	torch.tensor(lowercase_)
          self.assertTrue(np.allclose(reshape(lowercase_		,			(4, 3))		,			reshape(lowercase_		,			(4, 3)).numpy()))

          SCREAMING_SNAKE_CASE_   :						int							=	np.random.randn(3		,			4		,			5)
          SCREAMING_SNAKE_CASE_   :						int							=	torch.tensor(lowercase_)
          self.assertTrue(np.allclose(reshape(lowercase_		,			(12, 5))		,			reshape(lowercase_		,			(12, 5)).numpy()))
      @require_tf
      def _SCREAMING_SNAKE_CASE				(		self  :						Dict):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						int							=	tf.constant(lowercase_)
          self.assertTrue(np.allclose(reshape(lowercase_		,			(4, 3))		,			reshape(lowercase_		,			(4, 3)).numpy()))

          SCREAMING_SNAKE_CASE_   :						List[Any]							=	np.random.randn(3		,			4		,			5)
          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	tf.constant(lowercase_)
          self.assertTrue(np.allclose(reshape(lowercase_		,			(12, 5))		,			reshape(lowercase_		,			(12, 5)).numpy()))
      @require_flax
      def _SCREAMING_SNAKE_CASE				(		self  :						Any):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Any							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	jnp.array(lowercase_)
          self.assertTrue(np.allclose(reshape(lowercase_		,			(4, 3))		,			np.asarray(reshape(lowercase_		,			(4, 3)))))

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	np.random.randn(3		,			4		,			5)
          SCREAMING_SNAKE_CASE_   :						List[Any]							=	jnp.array(lowercase_)
          self.assertTrue(np.allclose(reshape(lowercase_		,			(12, 5))		,			np.asarray(reshape(lowercase_		,			(12, 5)))))
      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						List[Any]							=	np.random.randn(1		,			3		,			4)
          self.assertTrue(np.allclose(squeeze(lowercase_)		,			np.squeeze(lowercase_)))

          SCREAMING_SNAKE_CASE_   :						str							=	np.random.randn(1		,			4		,			1		,			5)
          self.assertTrue(np.allclose(squeeze(lowercase_		,			axis=2)		,			np.squeeze(lowercase_		,			axis=2)))
      @require_torch
      def _SCREAMING_SNAKE_CASE				(		self  :						Any):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	np.random.randn(1		,			3		,			4)
          SCREAMING_SNAKE_CASE_   :						int							=	torch.tensor(lowercase_)
          self.assertTrue(np.allclose(squeeze(lowercase_)		,			squeeze(lowercase_).numpy()))

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	np.random.randn(1		,			4		,			1		,			5)
          SCREAMING_SNAKE_CASE_   :						str							=	torch.tensor(lowercase_)
          self.assertTrue(np.allclose(squeeze(lowercase_		,			axis=2)		,			squeeze(lowercase_		,			axis=2).numpy()))
      @require_tf
      def _SCREAMING_SNAKE_CASE				(		self  :						List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						int							=	np.random.randn(1		,			3		,			4)
          SCREAMING_SNAKE_CASE_   :						Dict							=	tf.constant(lowercase_)
          self.assertTrue(np.allclose(squeeze(lowercase_)		,			squeeze(lowercase_).numpy()))

          SCREAMING_SNAKE_CASE_   :						List[Any]							=	np.random.randn(1		,			4		,			1		,			5)
          SCREAMING_SNAKE_CASE_   :						Any							=	tf.constant(lowercase_)
          self.assertTrue(np.allclose(squeeze(lowercase_		,			axis=2)		,			squeeze(lowercase_		,			axis=2).numpy()))
      @require_flax
      def _SCREAMING_SNAKE_CASE				(		self  :						int):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	np.random.randn(1		,			3		,			4)
          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	jnp.array(lowercase_)
          self.assertTrue(np.allclose(squeeze(lowercase_)		,			np.asarray(squeeze(lowercase_))))

          SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	np.random.randn(1		,			4		,			1		,			5)
          SCREAMING_SNAKE_CASE_   :						str							=	jnp.array(lowercase_)
          self.assertTrue(np.allclose(squeeze(lowercase_		,			axis=2)		,			np.asarray(squeeze(lowercase_		,			axis=2))))
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						int							=	np.random.randn(3		,			4)
          self.assertTrue(np.allclose(expand_dims(lowercase_		,			axis=1)		,			np.expand_dims(lowercase_		,			axis=1)))
      @require_torch
      def _SCREAMING_SNAKE_CASE				(		self  :						Dict):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						int							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						int							=	torch.tensor(lowercase_)
          self.assertTrue(np.allclose(expand_dims(lowercase_		,			axis=1)		,			expand_dims(lowercase_		,			axis=1).numpy()))
      @require_tf
      def _SCREAMING_SNAKE_CASE				(		self  :						str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						List[str]							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						Any							=	tf.constant(lowercase_)
          self.assertTrue(np.allclose(expand_dims(lowercase_		,			axis=1)		,			expand_dims(lowercase_		,			axis=1).numpy()))



      @require_flax
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Tuple							=	np.random.randn(3		,			4)
          SCREAMING_SNAKE_CASE_   :						int							=	jnp.array(lowercase_)
          self.assertTrue(np.allclose(expand_dims(lowercase_		,			axis=1)		,			np.asarray(expand_dims(lowercase_		,			axis=1))))


