


from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...file_utils import TensorType, is_torch_available
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import logging


SCREAMING_SNAKE_CASE       :      List[Any]												= logging.get_logger(__name__)

SCREAMING_SNAKE_CASE       :      Optional[Any]												= {
    "facebook/blenderbot_small-90M": "https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/config.json",
    # See all BlenderbotSmall models at https://huggingface.co/models?filter=blenderbot_small
}






class _lowerCamelCase(						_a   ):
       lowercase_ :	Any     =			"""blenderbot-small"""
       lowercase_ :	List[Any]     =			["""past_key_values"""]
       lowercase_ :	Union[str, Any]     =			{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}
       def __init__(					self,      lowerCamelCase=5_02_65,      lowerCamelCase=5_12,      lowerCamelCase=8,      lowerCamelCase=20_48,      lowerCamelCase=16,      lowerCamelCase=8,      lowerCamelCase=20_48,      lowerCamelCase=16,      lowerCamelCase=0.0,      lowerCamelCase=0.0,      lowerCamelCase=True,      lowerCamelCase=True,      lowerCamelCase="gelu",      lowerCamelCase=5_12,      lowerCamelCase=0.1,      lowerCamelCase=0.0,      lowerCamelCase=0.0,      lowerCamelCase=0.0_2,      lowerCamelCase=1,      lowerCamelCase=False,      lowerCamelCase=0,      lowerCamelCase=1,      lowerCamelCase=2,      lowerCamelCase=2,      **lowerCamelCase,      )       ->   List[str]:




           """simple docstring"""

           _lowercase       :							Dict      		=		vocab_size
           _lowercase       :							Dict      		=		max_position_embeddings
           _lowercase       :							Optional[int]      		=		d_model
           _lowercase       :							str      		=		encoder_ffn_dim
           _lowercase       :							str      		=		encoder_layers
           _lowercase       :							Any      		=		encoder_attention_heads
           _lowercase       :							Any      		=		decoder_ffn_dim
           _lowercase       :							List[str]      		=		decoder_layers
           _lowercase       :							str      		=		decoder_attention_heads
           _lowercase       :							Dict      		=		dropout
           _lowercase       :							str      		=		attention_dropout
           _lowercase       :							Union[str, Any]      		=		activation_dropout
           _lowercase       :							Tuple      		=		activation_function
           _lowercase       :							Optional[int]      		=		init_std
           _lowercase       :							Any      		=		encoder_layerdrop
           _lowercase       :							Optional[int]      		=		decoder_layerdrop
           _lowercase       :							Optional[Any]      		=		use_cache
           _lowercase       :							Dict      		=		encoder_layers
           _lowercase       :							List[str]      		=		scale_embedding  # scale factor will be sqrt(d_model) if True

           super().__init__(
               pad_token_id=lowerCamelCase,      bos_token_id=lowerCamelCase,      eos_token_id=lowerCamelCase,      is_encoder_decoder=lowerCamelCase,      decoder_start_token_id=lowerCamelCase,      forced_eos_token_id=lowerCamelCase,      **lowerCamelCase,      )






class _lowerCamelCase(						_a   ):
       @property
       def 			UpperCamelCase     (					self)       ->   Mapping[str, Mapping[int, str]]:




           """simple docstring"""

           if self.task in ["default", "seq2seq-lm"]:
               _lowercase       :							List[Any]      		=		OrderedDict(
                   [
                       ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
                       ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
                   ])

               if self.use_past:
                   _lowercase       :							Tuple      		=		{0: 'batch'}
                   _lowercase       :							Optional[int]      		=		{0: 'batch', 1: 'past_decoder_sequence + sequence'}
               else:
                   _lowercase       :							str      		=		{0: 'batch', 1: 'decoder_sequence'}
                   _lowercase       :							int      		=		{0: 'batch', 1: 'decoder_sequence'}

               if self.use_past:
                   self.fill_with_past_key_values_(lowerCamelCase,      direction='inputs')
           elif self.task == "causal-lm":
               # TODO: figure this case out.
               _lowercase       :							Tuple      		=		OrderedDict(
                   [
                       ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
                       ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
                   ])
               if self.use_past:
                   _lowercase		,				_lowercase       :							List[str]      		=		self.num_layers
                   for i in range(lowerCamelCase):
                       _lowercase       :							int      		=		{0: 'batch', 2: 'past_sequence + sequence'}
                       _lowercase       :							str      		=		{0: 'batch', 2: 'past_sequence + sequence'}
           else:
               _lowercase       :							str      		=		OrderedDict(
                   [
                       ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
                       ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
                       ('decoder_input_ids', {0: 'batch', 1: 'decoder_sequence'}),
                       ('decoder_attention_mask', {0: 'batch', 1: 'decoder_sequence'}),
                   ])

           return common_inputs
       @property
       def 			UpperCamelCase     (					self)       ->   Mapping[str, Mapping[int, str]]:




           """simple docstring"""

           if self.task in ["default", "seq2seq-lm"]:
               _lowercase       :							Any      		=		super().outputs
           else:
               _lowercase       :							Tuple      		=		super(lowerCamelCase,      self).outputs
               if self.use_past:
                   _lowercase		,				_lowercase       :							int      		=		self.num_layers
                   for i in range(lowerCamelCase):
                       _lowercase       :							Any      		=		{0: 'batch', 2: 'past_sequence + sequence'}
                       _lowercase       :							Tuple      		=		{0: 'batch', 2: 'past_sequence + sequence'}
           return common_outputs
       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase = -1,      lowerCamelCase = -1,      lowerCamelCase = False,      lowerCamelCase = None,      )       ->   Mapping[str, Any]:




           """simple docstring"""

           _lowercase       :							Optional[int]      		=		self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
               lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)

           # Generate decoder inputs
           _lowercase       :							str      		=		seq_length if not self.use_past else 1
           _lowercase       :							List[Any]      		=		self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
               lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)
           _lowercase       :							Optional[int]      		=		{F'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
           _lowercase       :							Any      		=		dict(**lowerCamelCase,      **lowerCamelCase)

           if self.use_past:
               if not is_torch_available():
                   raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.')
               else:
                   import torch
               _lowercase		,				_lowercase       :							int      		=		common_inputs['input_ids'].shape
               _lowercase       :							Dict      		=		common_inputs['decoder_input_ids'].shape[1]
               _lowercase		,				_lowercase       :							int      		=		self.num_attention_heads
               _lowercase       :							Tuple      		=		(
                   batch,
                   num_encoder_attention_heads,
                   encoder_seq_length,
                   self._config.hidden_size // num_encoder_attention_heads,
               )
               _lowercase       :							Tuple      		=		decoder_seq_length + 3
               _lowercase       :							str      		=		(
                   batch,
                   num_decoder_attention_heads,
                   decoder_past_length,
                   self._config.hidden_size // num_decoder_attention_heads,
               )

               _lowercase       :							int      		=		torch.cat(
                   [common_inputs['decoder_attention_mask'], torch.ones(lowerCamelCase,      lowerCamelCase)],      dim=1)

               _lowercase       :							Optional[int]      		=		[]
               # If the number of encoder and decoder layers are present in the model configuration, both are considered
               _lowercase		,				_lowercase       :							Dict      		=		self.num_layers
               _lowercase       :							List[str]      		=		min(lowerCamelCase,      lowerCamelCase)
               _lowercase       :							str      		=		max(lowerCamelCase,      lowerCamelCase) - min_num_layers
               _lowercase       :							Optional[int]      		=		'encoder' if num_encoder_layers > num_decoder_layers else 'decoder'

               for _ in range(lowerCamelCase):
                   common_inputs["past_key_values"].append(
                       (
                           torch.zeros(lowerCamelCase),
                           torch.zeros(lowerCamelCase),
                           torch.zeros(lowerCamelCase),
                           torch.zeros(lowerCamelCase),
                       ))
               # TODO: test this.
               _lowercase       :							List[Any]      		=		encoder_shape if remaining_side_name == 'encoder' else decoder_shape
               for _ in range(lowerCamelCase,      lowerCamelCase):
                   common_inputs["past_key_values"].append((torch.zeros(lowerCamelCase), torch.zeros(lowerCamelCase)))
           return common_inputs
       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase = -1,      lowerCamelCase = -1,      lowerCamelCase = False,      lowerCamelCase = None,      )       ->   Mapping[str, Any]:




           """simple docstring"""

           _lowercase       :							Any      		=		self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
               lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)

           if self.use_past:
               if not is_torch_available():
                   raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.')
               else:
                   import torch
               _lowercase		,				_lowercase       :							Optional[Any]      		=		common_inputs['input_ids'].shape
               # Not using the same length for past_key_values
               _lowercase       :							Any      		=		seqlen + 2
               _lowercase		,				_lowercase       :							Dict      		=		self.num_layers
               _lowercase		,				_lowercase       :							Optional[Any]      		=		self.num_attention_heads
               _lowercase       :							Union[str, Any]      		=		(
                   batch,
                   num_encoder_attention_heads,
                   past_key_values_length,
                   self._config.hidden_size // num_encoder_attention_heads,
               )

               _lowercase       :							Tuple      		=		common_inputs['attention_mask'].dtype
               _lowercase       :							Optional[int]      		=		torch.cat(
                   [common_inputs['attention_mask'], torch.ones(lowerCamelCase,      lowerCamelCase,      dtype=lowerCamelCase)],      dim=1)
               _lowercase       :							Any      		=		[
                   (torch.zeros(lowerCamelCase), torch.zeros(lowerCamelCase)) for _ in range(lowerCamelCase)
               ]
           return common_inputs
       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase = -1,      lowerCamelCase = -1,      lowerCamelCase = False,      lowerCamelCase = None,      )       ->   Mapping[str, Any]:




           """simple docstring"""

           _lowercase       :							Tuple      		=		compute_effective_axis_dimension(
               lowerCamelCase,      fixed_dimension=OnnxConfig.default_fixed_batch,      num_token_to_add=0)

           # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
           _lowercase       :							List[str]      		=		tokenizer.num_special_tokens_to_add(lowerCamelCase)
           _lowercase       :							int      		=		compute_effective_axis_dimension(
               lowerCamelCase,      fixed_dimension=OnnxConfig.default_fixed_sequence,      num_token_to_add=lowerCamelCase)

           # Generate dummy inputs according to compute batch and sequence
           _lowercase       :							Optional[int]      		=		[' '.join([tokenizer.unk_token]) * seq_length] * batch_size
           _lowercase       :							List[str]      		=		dict(tokenizer(lowerCamelCase,      return_tensors=lowerCamelCase))
           return common_inputs
       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase = -1,      lowerCamelCase = -1,      lowerCamelCase = False,      lowerCamelCase = None,      )       ->   Mapping[str, Any]:




           """simple docstring"""

           if self.task in ["default", "seq2seq-lm"]:
               _lowercase       :							Optional[Any]      		=		self._generate_dummy_inputs_for_default_and_seqaseq_lm(
                   lowerCamelCase,      batch_size=lowerCamelCase,      seq_length=lowerCamelCase,      is_pair=lowerCamelCase,      framework=lowerCamelCase)

           elif self.task == "causal-lm":
               _lowercase       :							Any      		=		self._generate_dummy_inputs_for_causal_lm(
                   lowerCamelCase,      batch_size=lowerCamelCase,      seq_length=lowerCamelCase,      is_pair=lowerCamelCase,      framework=lowerCamelCase)
           else:
               _lowercase       :							Optional[int]      		=		self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
                   lowerCamelCase,      batch_size=lowerCamelCase,      seq_length=lowerCamelCase,      is_pair=lowerCamelCase,      framework=lowerCamelCase)

           return common_inputs




       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)       ->   int:




           """simple docstring"""

           if self.task in ["default", "seq2seq-lm"]:
               _lowercase       :							List[Any]      		=		super()._flatten_past_key_values_(lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)
           else:
               _lowercase       :							Tuple      		=		super(lowerCamelCase,      self)._flatten_past_key_values_(
                   lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)








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       lowercase_    (      self							)    ->      Tuple:



							'''simple docstring'''
							__lowerCamelCase		       =    {
							    '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},
							    }
							}
							__lowerCamelCase		       =    {
							    '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(lowerCamelCase__							)			,						lowerCamelCase__							)


			def       lowercase_    (      self							)    ->      Tuple:



							'''simple docstring'''
							__lowerCamelCase		       =    np.random.randn(3			,						4							)
							self.assertTrue(np.allclose(transpose(lowerCamelCase__							)			,						x.transpose()							)							)

							__lowerCamelCase		       =    np.random.randn(3			,						4			,						5							)
							self.assertTrue(np.allclose(transpose(lowerCamelCase__			,						axes=(1, 2, 0)							)			,						x.transpose((1, 2, 0)							)							)							)


			@require_torch
			def       lowercase_    (      self							)    ->      int:



							'''simple docstring'''
							__lowerCamelCase		       =    np.random.randn(3			,						4							)
							__lowerCamelCase		       =    torch.tensor(lowerCamelCase__							)
							self.assertTrue(np.allclose(transpose(lowerCamelCase__							)			,						transpose(lowerCamelCase__							).numpy()							)							)

							__lowerCamelCase		       =    np.random.randn(3			,						4			,						5							)
							__lowerCamelCase		       =    torch.tensor(lowerCamelCase__							)
							self.assertTrue(np.allclose(transpose(lowerCamelCase__			,						axes=(1, 2, 0)							)			,						transpose(lowerCamelCase__			,						axes=(1, 2, 0)							).numpy()							)							)


			@require_tf
			def       lowercase_    (      self							)    ->      Union[str, Any]:



							'''simple docstring'''
							__lowerCamelCase		       =    np.random.randn(3			,						4							)
							__lowerCamelCase		       =    tf.constant(lowerCamelCase__							)
							self.assertTrue(np.allclose(transpose(lowerCamelCase__							)			,						transpose(lowerCamelCase__							).numpy()							)							)

							__lowerCamelCase		       =    np.random.randn(3			,						4			,						5							)
							__lowerCamelCase		       =    tf.constant(lowerCamelCase__							)
							self.assertTrue(np.allclose(transpose(lowerCamelCase__			,						axes=(1, 2, 0)							)			,						transpose(lowerCamelCase__			,						axes=(1, 2, 0)							).numpy()							)							)


			@require_flax
			def       lowercase_    (      self							)    ->      Optional[int]:



							'''simple docstring'''
							__lowerCamelCase		       =    np.random.randn(3			,						4							)
							__lowerCamelCase		       =    jnp.array(lowerCamelCase__							)
							self.assertTrue(np.allclose(transpose(lowerCamelCase__							)			,						np.asarray(transpose(lowerCamelCase__							)							)							)							)

							__lowerCamelCase		       =    np.random.randn(3			,						4			,						5							)
							__lowerCamelCase		       =    jnp.array(lowerCamelCase__							)
							self.assertTrue(np.allclose(transpose(lowerCamelCase__			,						axes=(1, 2, 0)							)			,						np.asarray(transpose(lowerCamelCase__			,						axes=(1, 2, 0)							)							)							)							)


			def       lowercase_    (      self							)    ->      int:



							'''simple docstring'''
							__lowerCamelCase		       =    np.random.randn(3			,						4							)
							self.assertTrue(np.allclose(reshape(lowerCamelCase__			,						(4, 3)							)			,						np.reshape(lowerCamelCase__			,						(4, 3)							)							)							)

							__lowerCamelCase		       =    np.random.randn(3			,						4			,						5							)
							self.assertTrue(np.allclose(reshape(lowerCamelCase__			,						(12, 5)							)			,						np.reshape(lowerCamelCase__			,						(12, 5)							)							)							)


			@require_torch
			def       lowercase_    (      self							)    ->      List[Any]:



							'''simple docstring'''
							__lowerCamelCase		       =    np.random.randn(3			,						4							)
							__lowerCamelCase		       =    torch.tensor(lowerCamelCase__							)
							self.assertTrue(np.allclose(reshape(lowerCamelCase__			,						(4, 3)							)			,						reshape(lowerCamelCase__			,						(4, 3)							).numpy()							)							)

							__lowerCamelCase		       =    np.random.randn(3			,						4			,						5							)
							__lowerCamelCase		       =    torch.tensor(lowerCamelCase__							)
							self.assertTrue(np.allclose(reshape(lowerCamelCase__			,						(12, 5)							)			,						reshape(lowerCamelCase__			,						(12, 5)							).numpy()							)							)


			@require_tf
			def       lowercase_    (      self							)    ->      Optional[int]:



							'''simple docstring'''
							__lowerCamelCase		       =    np.random.randn(3			,						4							)
							__lowerCamelCase		       =    tf.constant(lowerCamelCase__							)
							self.assertTrue(np.allclose(reshape(lowerCamelCase__			,						(4, 3)							)			,						reshape(lowerCamelCase__			,						(4, 3)							).numpy()							)							)

							__lowerCamelCase		       =    np.random.randn(3			,						4			,						5							)
							__lowerCamelCase		       =    tf.constant(lowerCamelCase__							)
							self.assertTrue(np.allclose(reshape(lowerCamelCase__			,						(12, 5)							)			,						reshape(lowerCamelCase__			,						(12, 5)							).numpy()							)							)


			@require_flax
			def       lowercase_    (      self							)    ->      Tuple:



							'''simple docstring'''
							__lowerCamelCase		       =    np.random.randn(3			,						4							)
							__lowerCamelCase		       =    jnp.array(lowerCamelCase__							)
							self.assertTrue(np.allclose(reshape(lowerCamelCase__			,						(4, 3)							)			,						np.asarray(reshape(lowerCamelCase__			,						(4, 3)							)							)							)							)

							__lowerCamelCase		       =    np.random.randn(3			,						4			,						5							)
							__lowerCamelCase		       =    jnp.array(lowerCamelCase__							)
							self.assertTrue(np.allclose(reshape(lowerCamelCase__			,						(12, 5)							)			,						np.asarray(reshape(lowerCamelCase__			,						(12, 5)							)							)							)							)


			def       lowercase_    (      self							)    ->      int:



							'''simple docstring'''
							__lowerCamelCase		       =    np.random.randn(1			,						3			,						4							)
							self.assertTrue(np.allclose(squeeze(lowerCamelCase__							)			,						np.squeeze(lowerCamelCase__							)							)							)

							__lowerCamelCase		       =    np.random.randn(1			,						4			,						1			,						5							)
							self.assertTrue(np.allclose(squeeze(lowerCamelCase__			,						axis=2							)			,						np.squeeze(lowerCamelCase__			,						axis=2							)							)							)


			@require_torch
			def       lowercase_    (      self							)    ->      int:



							'''simple docstring'''
							__lowerCamelCase		       =    np.random.randn(1			,						3			,						4							)
							__lowerCamelCase		       =    torch.tensor(lowerCamelCase__							)
							self.assertTrue(np.allclose(squeeze(lowerCamelCase__							)			,						squeeze(lowerCamelCase__							).numpy()							)							)

							__lowerCamelCase		       =    np.random.randn(1			,						4			,						1			,						5							)
							__lowerCamelCase		       =    torch.tensor(lowerCamelCase__							)
							self.assertTrue(np.allclose(squeeze(lowerCamelCase__			,						axis=2							)			,						squeeze(lowerCamelCase__			,						axis=2							).numpy()							)							)


			@require_tf
			def       lowercase_    (      self							)    ->      Union[str, Any]:



							'''simple docstring'''
							__lowerCamelCase		       =    np.random.randn(1			,						3			,						4							)
							__lowerCamelCase		       =    tf.constant(lowerCamelCase__							)
							self.assertTrue(np.allclose(squeeze(lowerCamelCase__							)			,						squeeze(lowerCamelCase__							).numpy()							)							)

							__lowerCamelCase		       =    np.random.randn(1			,						4			,						1			,						5							)
							__lowerCamelCase		       =    tf.constant(lowerCamelCase__							)
							self.assertTrue(np.allclose(squeeze(lowerCamelCase__			,						axis=2							)			,						squeeze(lowerCamelCase__			,						axis=2							).numpy()							)							)


			@require_flax
			def       lowercase_    (      self							)    ->      List[str]:



							'''simple docstring'''
							__lowerCamelCase		       =    np.random.randn(1			,						3			,						4							)
							__lowerCamelCase		       =    jnp.array(lowerCamelCase__							)
							self.assertTrue(np.allclose(squeeze(lowerCamelCase__							)			,						np.asarray(squeeze(lowerCamelCase__							)							)							)							)

							__lowerCamelCase		       =    np.random.randn(1			,						4			,						1			,						5							)
							__lowerCamelCase		       =    jnp.array(lowerCamelCase__							)
							self.assertTrue(np.allclose(squeeze(lowerCamelCase__			,						axis=2							)			,						np.asarray(squeeze(lowerCamelCase__			,						axis=2							)							)							)							)


			def       lowercase_    (      self							)    ->      int:



							'''simple docstring'''
							__lowerCamelCase		       =    np.random.randn(3			,						4							)
							self.assertTrue(np.allclose(expand_dims(lowerCamelCase__			,						axis=1							)			,						np.expand_dims(lowerCamelCase__			,						axis=1							)							)							)


			@require_torch
			def       lowercase_    (      self							)    ->      Any:



							'''simple docstring'''
							__lowerCamelCase		       =    np.random.randn(3			,						4							)
							__lowerCamelCase		       =    torch.tensor(lowerCamelCase__							)
							self.assertTrue(np.allclose(expand_dims(lowerCamelCase__			,						axis=1							)			,						expand_dims(lowerCamelCase__			,						axis=1							).numpy()							)							)


			@require_tf
			def       lowercase_    (      self							)    ->      Any:



							'''simple docstring'''
							__lowerCamelCase		       =    np.random.randn(3			,						4							)
							__lowerCamelCase		       =    tf.constant(lowerCamelCase__							)
							self.assertTrue(np.allclose(expand_dims(lowerCamelCase__			,						axis=1							)			,						expand_dims(lowerCamelCase__			,						axis=1							).numpy()							)							)






			@require_flax
			def       lowercase_    (      self							)    ->      List[Any]:



							'''simple docstring'''
							__lowerCamelCase		       =    np.random.randn(3			,						4							)
							__lowerCamelCase		       =    jnp.array(lowerCamelCase__							)
							self.assertTrue(np.allclose(expand_dims(lowerCamelCase__			,						axis=1							)			,						np.asarray(expand_dims(lowerCamelCase__			,						axis=1							)							)							)							)


