




'''simple docstring'''


import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
				import os

				import jax.numpy as jnp
				from jax import jit

				from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
				from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

				_snake_case				:				str       							=		'0.12'  # assumed parallelism: 8


if is_torch_available():
				import torch
def 		snake_case_   (UpperCamelCase	:  List[str]			,						UpperCamelCase	:  Tuple			,						UpperCamelCase	:  str=None				):





							'''simple docstring'''


							if rng is None:
														_a     					=    random.Random()

							_a     					=    1
							for dim in shape:
														total_dims *= dim

							_a     					=    []
							for _ in range(UpperCamelCase				):
														values.append(rng.randint(0			,						vocab_size - 1				)				)

							_a     					=    np.array(UpperCamelCase			,						dtype=jnp.intaa				).reshape(UpperCamelCase				)

							return output
def 		snake_case_   (UpperCamelCase	:  List[str]			,						UpperCamelCase	:  str=None				):





							'''simple docstring'''


							_a     					=    ids_tensor(UpperCamelCase			,						vocab_size=2			,						rng=UpperCamelCase				)
							# make sure that at least one token is attended to for each batch
							_a     					=    1
							return attn_mask

@require_flax
class       A   :
							lowercase_          =    None
							lowercase_          =    ()







							def        __lowerCAmelCase      (  self	:  str )      ->							Union[str, Any]:


														"""simple docstring"""


														_a					,							_a     					=    self.model_tester.prepare_config_and_inputs_for_common()

														# cut to half length & take max batch_size 3
														_a     					=    2
														_a     					=    inputs['''input_ids'''].shape[-1] // 2
														_a     					=    inputs['''input_ids'''][:max_batch_size, :sequence_length]

														_a     					=    jnp.ones_like(lowerCAmelCase_ )
														_a     					=    attention_mask[:max_batch_size, :sequence_length]

														# generate max 5 tokens
														_a     					=    input_ids.shape[-1] + 5
														if config.eos_token_id is not None and config.pad_token_id is None:
																					# hack to allow generate for models such as GPT2 as is done in `generate()`
																					_a     					=    config.eos_token_id
														return config, input_ids, attention_mask, max_length







							@is_pt_flax_cross_test
							def        __lowerCAmelCase      (  self	:  Optional[Any] )      ->							str:


														"""simple docstring"""


														_a					,							_a					,							_a					,							_a     					=    self._get_input_ids_and_config()
														_a     					=    False
														_a     					=    max_length
														_a     					=    0

														for model_class in self.all_generative_model_classes:
																					_a     					=    model_class(lowerCAmelCase_ )

																					_a     					=    model_class.__name__[4:]  # Skip the "Flax" at the beginning
																					_a     					=    getattr(lowerCAmelCase_							,      lowerCAmelCase_ )
																					_a     					=    pt_model_class(lowerCAmelCase_ ).eval()
																					_a     					=    load_flax_weights_in_pytorch_model(lowerCAmelCase_							,      flax_model.params )

																					_a     					=    flax_model.generate(lowerCAmelCase_ ).sequences
																					_a     					=    pt_model.generate(torch.tensor(lowerCAmelCase_							,      dtype=torch.long ) )

																					if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
																												_a     					=    flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

																					self.assertListEqual(pt_generation_outputs.numpy().tolist()							,      flax_generation_outputs.tolist() )







							def        __lowerCAmelCase      (  self	:  List[str] )      ->							List[str]:


														"""simple docstring"""


														_a					,							_a					,							_a					,							_a     					=    self._get_input_ids_and_config()
														_a     					=    False
														_a     					=    max_length

														for model_class in self.all_generative_model_classes:
																					_a     					=    model_class(lowerCAmelCase_ )

																					_a     					=    model.generate(lowerCAmelCase_ ).sequences
																					self.assertEqual(generation_outputs.shape[-1]							,      lowerCAmelCase_ )

																					_a     					=    jit(model.generate )
																					_a     					=    jit_generate(lowerCAmelCase_ ).sequences

																					self.assertListEqual(generation_outputs.tolist()							,      jit_generation_outputs.tolist() )







							def        __lowerCAmelCase      (  self	:  List[Any] )      ->							Tuple:


														"""simple docstring"""


														_a					,							_a					,							_a					,							_a     					=    self._get_input_ids_and_config()
														_a     					=    True
														_a     					=    max_length

														for model_class in self.all_generative_model_classes:
																					_a     					=    model_class(lowerCAmelCase_ )

																					_a     					=    model.generate(lowerCAmelCase_ ).sequences
																					self.assertEqual(generation_outputs.shape[-1]							,      lowerCAmelCase_ )

																					_a     					=    jit(model.generate )
																					_a     					=    jit_generate(lowerCAmelCase_ ).sequences

																					self.assertListEqual(generation_outputs.tolist()							,      jit_generation_outputs.tolist() )







							def        __lowerCAmelCase      (  self	:  Union[str, Any] )      ->							Optional[Any]:


														"""simple docstring"""


														_a					,							_a					,							_a					,							_a     					=    self._get_input_ids_and_config()
														_a     					=    False
														_a     					=    max_length
														_a     					=    2

														for model_class in self.all_generative_model_classes:
																					_a     					=    model_class(lowerCAmelCase_ )

																					_a     					=    model.generate(lowerCAmelCase_ ).sequences
																					self.assertEqual(generation_outputs.shape[-1]							,      lowerCAmelCase_ )

																					_a     					=    jit(model.generate )
																					_a     					=    jit_generate(lowerCAmelCase_ ).sequences

																					self.assertListEqual(generation_outputs.tolist()							,      jit_generation_outputs.tolist() )







							def        __lowerCAmelCase      (  self	:  Tuple )      ->							Optional[Any]:


														"""simple docstring"""


														_a					,							_a					,							_a					,							_a     					=    self._get_input_ids_and_config()
														_a     					=    False
														_a     					=    max_length
														_a     					=    2
														_a     					=    2

														for model_class in self.all_generative_model_classes:
																					_a     					=    model_class(lowerCAmelCase_ )

																					_a     					=    model.generate(lowerCAmelCase_ ).sequences
																					self.assertEqual(generation_outputs.shape[0]							,      input_ids.shape[0] * config.num_return_sequences )







							def        __lowerCAmelCase      (  self	:  Optional[Any] )      ->							Union[str, Any]:


														"""simple docstring"""


														_a					,							_a					,							_a					,							_a     					=    self._get_input_ids_and_config()
														_a     					=    True
														_a     					=    max_length
														_a     					=    0.8
														_a     					=    10
														_a     					=    0.3
														_a     					=    1
														_a     					=    8
														_a     					=    9

														for model_class in self.all_generative_model_classes:
																					_a     					=    model_class(lowerCAmelCase_ )

																					_a     					=    model.generate(lowerCAmelCase_ ).sequences
																					self.assertEqual(generation_outputs.shape[-1]							,      lowerCAmelCase_ )

																					_a     					=    jit(model.generate )
																					_a     					=    jit_generate(lowerCAmelCase_ ).sequences

																					self.assertListEqual(generation_outputs.tolist()							,      jit_generation_outputs.tolist() )







							def        __lowerCAmelCase      (  self	:  Any )      ->							Optional[Any]:


														"""simple docstring"""


														_a					,							_a					,							_a					,							_a     					=    self._get_input_ids_and_config()
														_a     					=    max_length
														_a     					=    1
														_a     					=    8
														_a     					=    9

														for model_class in self.all_generative_model_classes:
																					_a     					=    model_class(lowerCAmelCase_ )

																					_a     					=    model.generate(lowerCAmelCase_ ).sequences
																					self.assertEqual(generation_outputs.shape[-1]							,      lowerCAmelCase_ )

																					_a     					=    jit(model.generate )
																					_a     					=    jit_generate(lowerCAmelCase_ ).sequences

																					self.assertListEqual(generation_outputs.tolist()							,      jit_generation_outputs.tolist() )







							def        __lowerCAmelCase      (  self	:  Optional[int] )      ->							str:


														"""simple docstring"""


														_a					,							_a					,							_a					,							_a     					=    self._get_input_ids_and_config()
														_a     					=    max_length
														_a     					=    2
														_a     					=    1
														_a     					=    8
														_a     					=    9

														for model_class in self.all_generative_model_classes:
																					_a     					=    model_class(lowerCAmelCase_ )

																					_a     					=    model.generate(lowerCAmelCase_ ).sequences
																					self.assertEqual(generation_outputs.shape[-1]							,      lowerCAmelCase_ )

																					_a     					=    jit(model.generate )
																					_a     					=    jit_generate(lowerCAmelCase_ ).sequences

																					self.assertListEqual(generation_outputs.tolist()							,      jit_generation_outputs.tolist() )







							def        __lowerCAmelCase      (  self	:  Any )      ->							Optional[int]:


														"""simple docstring"""


														_a					,							_a					,							_a					,							_a     					=    self._get_input_ids_and_config()

														# pad attention mask on the left
														_a     					=    attention_mask.at[(0, 0)].set(0 )

														_a     					=    False
														_a     					=    max_length

														for model_class in self.all_generative_model_classes:
																					_a     					=    model_class(lowerCAmelCase_ )

																					_a     					=    model.generate(lowerCAmelCase_							,      attention_mask=lowerCAmelCase_ ).sequences
																					self.assertEqual(generation_outputs.shape[-1]							,      lowerCAmelCase_ )

																					_a     					=    jit(model.generate )
																					_a     					=    jit_generate(lowerCAmelCase_							,      attention_mask=lowerCAmelCase_ ).sequences

																					self.assertListEqual(generation_outputs.tolist()							,      jit_generation_outputs.tolist() )







							def        __lowerCAmelCase      (  self	:  Optional[int] )      ->							str:


														"""simple docstring"""


														_a					,							_a					,							_a					,							_a     					=    self._get_input_ids_and_config()

														# pad attention mask on the left
														_a     					=    attention_mask.at[(0, 0)].set(0 )

														_a     					=    True
														_a     					=    max_length

														for model_class in self.all_generative_model_classes:
																					_a     					=    model_class(lowerCAmelCase_ )

																					_a     					=    model.generate(lowerCAmelCase_							,      attention_mask=lowerCAmelCase_ ).sequences
																					self.assertEqual(generation_outputs.shape[-1]							,      lowerCAmelCase_ )

																					_a     					=    jit(model.generate )
																					_a     					=    jit_generate(lowerCAmelCase_							,      attention_mask=lowerCAmelCase_ ).sequences

																					self.assertListEqual(generation_outputs.tolist()							,      jit_generation_outputs.tolist() )







							def        __lowerCAmelCase      (  self	:  Dict )      ->							Any:


														"""simple docstring"""


														_a					,							_a					,							_a					,							_a     					=    self._get_input_ids_and_config()

														# pad attention mask on the left
														_a     					=    attention_mask.at[(0, 0)].set(0 )

														_a     					=    2
														_a     					=    max_length

														for model_class in self.all_generative_model_classes:
																					_a     					=    model_class(lowerCAmelCase_ )

																					_a     					=    model.generate(lowerCAmelCase_							,      attention_mask=lowerCAmelCase_ ).sequences
																					self.assertEqual(generation_outputs.shape[-1]							,      lowerCAmelCase_ )

																					_a     					=    jit(model.generate )
																					_a     					=    jit_generate(lowerCAmelCase_							,      attention_mask=lowerCAmelCase_ ).sequences

																					self.assertListEqual(generation_outputs.tolist()							,      jit_generation_outputs.tolist() )



@require_flax
class       A   (		unittest.TestCase   ):







							def        __lowerCAmelCase      (  self	:  List[str] )      ->							Any:


														"""simple docstring"""


														_a     					=    AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-bert''' )
														_a     					=    FlaxAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-bert-flax-only''' )

														_a     					=    '''Hello world'''
														_a     					=    tokenizer(lowerCAmelCase_							,      return_tensors='''np''' ).input_ids

														# typos are quickly detected (the correct argument is `do_sample`)
														with self.assertRaisesRegex(lowerCAmelCase_							,      '''do_samples''' ):
																					model.generate(lowerCAmelCase_							,      do_samples=lowerCAmelCase_ )

														# arbitrary arguments that will not be used anywhere are also not accepted
														with self.assertRaisesRegex(lowerCAmelCase_							,      '''foo''' ):
																					_a     					=    {'''foo''': '''bar'''}
																					model.generate(lowerCAmelCase_							,      **lowerCAmelCase_ )







'''simple docstring'''


from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


_snake_case				:				Dict       							=		logging.get_logger(__name__)

_snake_case				:				Optional[Any]       							=		{
    'shi-labs/nat-mini-in1k-224': 'https://huggingface.co/shi-labs/nat-mini-in1k-224/resolve/main/config.json',
    # See all Nat models at https://huggingface.co/models?filter=nat
}

class       A   (		_a					,_a   ):
							lowercase_          =    'nat'

							lowercase_          =    {
							    'num_attention_heads': 'num_heads',
							    'num_hidden_layers': 'num_layers',
							}







							def __init__(  self	:  List[str]							,      lowerCAmelCase_	:  str=4							,      lowerCAmelCase_	:  Any=3							,      lowerCAmelCase_	:  List[Any]=64							,      lowerCAmelCase_	:  Dict=[3, 4, 6, 5]							,      lowerCAmelCase_	:  Dict=[2, 4, 8, 16]							,      lowerCAmelCase_	:  str=7							,      lowerCAmelCase_	:  Dict=3.0							,      lowerCAmelCase_	:  int=True							,      lowerCAmelCase_	:  Tuple=0.0							,      lowerCAmelCase_	:  str=0.0							,      lowerCAmelCase_	:  List[str]=0.1							,      lowerCAmelCase_	:  int="gelu"							,      lowerCAmelCase_	:  List[str]=0.0_2							,      lowerCAmelCase_	:  str=1e-5							,      lowerCAmelCase_	:  Tuple=0.0							,      lowerCAmelCase_	:  str=None							,      lowerCAmelCase_	:  int=None							,      **lowerCAmelCase_	:  List[Any]							,      )      ->							Any:


														"""simple docstring"""


														super().__init__(**lowerCAmelCase_ )

														_a     					=    patch_size
														_a     					=    num_channels
														_a     					=    embed_dim
														_a     					=    depths
														_a     					=    len(lowerCAmelCase_ )
														_a     					=    num_heads
														_a     					=    kernel_size
														_a     					=    mlp_ratio
														_a     					=    qkv_bias
														_a     					=    hidden_dropout_prob
														_a     					=    attention_probs_dropout_prob
														_a     					=    drop_path_rate
														_a     					=    hidden_act
														_a     					=    layer_norm_eps
														_a     					=    initializer_range
														# we set the hidden_size attribute in order to make Nat work with VisionEncoderDecoderModel
														# this indicates the channel dimension after the last stage of the model
														_a     					=    int(embed_dim * 2 ** (len(lowerCAmelCase_ ) - 1) )
														_a     					=    layer_scale_init_value
														_a     					=    ['''stem'''] + [F'stage{idx}' for idx in range(1							,      len(lowerCAmelCase_ ) + 1 )]
														_a					,							_a     					=    get_aligned_output_features_output_indices(
														    out_features=lowerCAmelCase_							,      out_indices=lowerCAmelCase_							,      stage_names=self.stage_names )

