


'''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

    __lowerCamelCase												=						'''0.12'''  # assumed parallelism: 8


if is_torch_available():
    import torch


def        a__     (					UpperCamelCase_						:      Optional[Any],   UpperCamelCase_						:      Any,   UpperCamelCase_						:      int=None						):
   if rng is None:
      UpperCAmelCase__	:Optional[Any]      =      random.Random()

   UpperCAmelCase__	:int      =      1
   for dim in shape:
      total_dims *= dim

   UpperCAmelCase__	:List[str]      =      []
   for _ in range(UpperCamelCase_						):
      values.append(rng.randint(0,   vocab_size - 1						)						)

   UpperCAmelCase__	:Union[str, Any]      =      np.array(UpperCamelCase_,   dtype=jnp.intaa						).reshape(UpperCamelCase_						)

   return output


def        a__     (					UpperCamelCase_						:      List[str],   UpperCamelCase_						:      Any=None						):
   UpperCAmelCase__	:Union[str, Any]      =      ids_tensor(UpperCamelCase_,   vocab_size=2,   rng=UpperCamelCase_						)
   # make sure that at least one token is attended to for each batch
   UpperCAmelCase__	:Optional[int]      =      1
   return attn_mask




@require_flax
class 				UpperCAmelCase			:
 UpperCAmelCase				=					None
 UpperCAmelCase				=					()


 def        __SCREAMING_SNAKE_CASE					(      self							:	Any ):
    UpperCAmelCase__   ,						UpperCAmelCase__	:Any      =      self.model_tester.prepare_config_and_inputs_for_common()

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

    UpperCAmelCase__	:Optional[Any]      =      jnp.ones_like(__lowerCamelCase )
    UpperCAmelCase__	:Optional[int]      =      attention_mask[:max_batch_size, :sequence_length]

    # generate max 5 tokens
    UpperCAmelCase__	:List[str]      =      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()`
       UpperCAmelCase__	:Optional[Any]      =      config.eos_token_id
    return config, input_ids, attention_mask, max_length


 @is_pt_flax_cross_test
 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[int] ):
    UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__	:int      =      self._get_input_ids_and_config()
    UpperCAmelCase__	:str      =      False
    UpperCAmelCase__	:Any      =      max_length
    UpperCAmelCase__	:str      =      0

    for model_class in self.all_generative_model_classes:
       UpperCAmelCase__	:Dict      =      model_class(__lowerCamelCase )

       UpperCAmelCase__	:List[str]      =      model_class.__name__[4:]  # Skip the "Flax" at the beginning
       UpperCAmelCase__	:Tuple      =      getattr(__lowerCamelCase			,       __lowerCamelCase )
       UpperCAmelCase__	:List[Any]      =      pt_model_class(__lowerCamelCase ).eval()
       UpperCAmelCase__	:int      =      load_flax_weights_in_pytorch_model(__lowerCamelCase			,       flax_model.params )

       UpperCAmelCase__	:List[Any]      =      flax_model.generate(__lowerCamelCase ).sequences
       UpperCAmelCase__	:int      =      pt_model.generate(torch.tensor(__lowerCamelCase			,       dtype=torch.long ) )

       if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
          UpperCAmelCase__	:Optional[Any]      =      flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

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


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[int] ):
    UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__	:Union[str, Any]      =      self._get_input_ids_and_config()
    UpperCAmelCase__	:List[str]      =      False
    UpperCAmelCase__	:Dict      =      max_length

    for model_class in self.all_generative_model_classes:
       UpperCAmelCase__	:Tuple      =      model_class(__lowerCamelCase )

       UpperCAmelCase__	:List[str]      =      model.generate(__lowerCamelCase ).sequences
       self.assertEqual(generation_outputs.shape[-1]			,       __lowerCamelCase )

       UpperCAmelCase__	:List[str]      =      jit(model.generate )
       UpperCAmelCase__	:Union[str, Any]      =      jit_generate(__lowerCamelCase ).sequences

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


 def        __SCREAMING_SNAKE_CASE					(      self							:	str ):
    UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__	:List[str]      =      self._get_input_ids_and_config()
    UpperCAmelCase__	:Dict      =      True
    UpperCAmelCase__	:List[str]      =      max_length

    for model_class in self.all_generative_model_classes:
       UpperCAmelCase__	:List[str]      =      model_class(__lowerCamelCase )

       UpperCAmelCase__	:Union[str, Any]      =      model.generate(__lowerCamelCase ).sequences
       self.assertEqual(generation_outputs.shape[-1]			,       __lowerCamelCase )

       UpperCAmelCase__	:Tuple      =      jit(model.generate )
       UpperCAmelCase__	:Optional[Any]      =      jit_generate(__lowerCamelCase ).sequences

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


 def        __SCREAMING_SNAKE_CASE					(      self							:	Any ):
    UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__	:Optional[int]      =      self._get_input_ids_and_config()
    UpperCAmelCase__	:Tuple      =      False
    UpperCAmelCase__	:Union[str, Any]      =      max_length
    UpperCAmelCase__	:List[Any]      =      2

    for model_class in self.all_generative_model_classes:
       UpperCAmelCase__	:Tuple      =      model_class(__lowerCamelCase )

       UpperCAmelCase__	:str      =      model.generate(__lowerCamelCase ).sequences
       self.assertEqual(generation_outputs.shape[-1]			,       __lowerCamelCase )

       UpperCAmelCase__	:Any      =      jit(model.generate )
       UpperCAmelCase__	:Any      =      jit_generate(__lowerCamelCase ).sequences

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


 def        __SCREAMING_SNAKE_CASE					(      self							:	Tuple ):
    UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__	:Dict      =      self._get_input_ids_and_config()
    UpperCAmelCase__	:Union[str, Any]      =      False
    UpperCAmelCase__	:Tuple      =      max_length
    UpperCAmelCase__	:Union[str, Any]      =      2
    UpperCAmelCase__	:Union[str, Any]      =      2

    for model_class in self.all_generative_model_classes:
       UpperCAmelCase__	:str      =      model_class(__lowerCamelCase )

       UpperCAmelCase__	:List[Any]      =      model.generate(__lowerCamelCase ).sequences
       self.assertEqual(generation_outputs.shape[0]			,       input_ids.shape[0] * config.num_return_sequences )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Any ):
    UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__	:List[str]      =      self._get_input_ids_and_config()
    UpperCAmelCase__	:Any      =      True
    UpperCAmelCase__	:Union[str, Any]      =      max_length
    UpperCAmelCase__	:Optional[Any]      =      0.8
    UpperCAmelCase__	:List[str]      =      1_0
    UpperCAmelCase__	:Dict      =      0.3
    UpperCAmelCase__	:Any      =      1
    UpperCAmelCase__	:Dict      =      8
    UpperCAmelCase__	:List[str]      =      9

    for model_class in self.all_generative_model_classes:
       UpperCAmelCase__	:List[Any]      =      model_class(__lowerCamelCase )

       UpperCAmelCase__	:Tuple      =      model.generate(__lowerCamelCase ).sequences
       self.assertEqual(generation_outputs.shape[-1]			,       __lowerCamelCase )

       UpperCAmelCase__	:Tuple      =      jit(model.generate )
       UpperCAmelCase__	:Dict      =      jit_generate(__lowerCamelCase ).sequences

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


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[Any] ):
    UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__	:List[str]      =      self._get_input_ids_and_config()
    UpperCAmelCase__	:List[Any]      =      max_length
    UpperCAmelCase__	:List[str]      =      1
    UpperCAmelCase__	:Dict      =      8
    UpperCAmelCase__	:Tuple      =      9

    for model_class in self.all_generative_model_classes:
       UpperCAmelCase__	:Any      =      model_class(__lowerCamelCase )

       UpperCAmelCase__	:Any      =      model.generate(__lowerCamelCase ).sequences
       self.assertEqual(generation_outputs.shape[-1]			,       __lowerCamelCase )

       UpperCAmelCase__	:Any      =      jit(model.generate )
       UpperCAmelCase__	:Union[str, Any]      =      jit_generate(__lowerCamelCase ).sequences

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


 def        __SCREAMING_SNAKE_CASE					(      self							:	Tuple ):
    UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__	:Tuple      =      self._get_input_ids_and_config()
    UpperCAmelCase__	:str      =      max_length
    UpperCAmelCase__	:str      =      2
    UpperCAmelCase__	:List[str]      =      1
    UpperCAmelCase__	:List[str]      =      8
    UpperCAmelCase__	:str      =      9

    for model_class in self.all_generative_model_classes:
       UpperCAmelCase__	:Tuple      =      model_class(__lowerCamelCase )

       UpperCAmelCase__	:str      =      model.generate(__lowerCamelCase ).sequences
       self.assertEqual(generation_outputs.shape[-1]			,       __lowerCamelCase )

       UpperCAmelCase__	:str      =      jit(model.generate )
       UpperCAmelCase__	:List[Any]      =      jit_generate(__lowerCamelCase ).sequences

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


 def        __SCREAMING_SNAKE_CASE					(      self							:	Dict ):
    UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__	:Any      =      self._get_input_ids_and_config()

    # pad attention mask on the left
    UpperCAmelCase__	:List[str]      =      attention_mask.at[(0, 0)].set(0 )

    UpperCAmelCase__	:List[str]      =      False
    UpperCAmelCase__	:List[Any]      =      max_length

    for model_class in self.all_generative_model_classes:
       UpperCAmelCase__	:str      =      model_class(__lowerCamelCase )

       UpperCAmelCase__	:Optional[int]      =      model.generate(__lowerCamelCase			,       attention_mask=__lowerCamelCase ).sequences
       self.assertEqual(generation_outputs.shape[-1]			,       __lowerCamelCase )

       UpperCAmelCase__	:Any      =      jit(model.generate )
       UpperCAmelCase__	:Dict      =      jit_generate(__lowerCamelCase			,       attention_mask=__lowerCamelCase ).sequences

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


 def        __SCREAMING_SNAKE_CASE					(      self							:	int ):
    UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__	:Dict      =      self._get_input_ids_and_config()

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

    UpperCAmelCase__	:Optional[int]      =      True
    UpperCAmelCase__	:Tuple      =      max_length

    for model_class in self.all_generative_model_classes:
       UpperCAmelCase__	:Any      =      model_class(__lowerCamelCase )

       UpperCAmelCase__	:List[str]      =      model.generate(__lowerCamelCase			,       attention_mask=__lowerCamelCase ).sequences
       self.assertEqual(generation_outputs.shape[-1]			,       __lowerCamelCase )

       UpperCAmelCase__	:Optional[Any]      =      jit(model.generate )
       UpperCAmelCase__	:Tuple      =      jit_generate(__lowerCamelCase			,       attention_mask=__lowerCamelCase ).sequences

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





 def        __SCREAMING_SNAKE_CASE					(      self							:	List[Any] ):
    UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__	:Dict      =      self._get_input_ids_and_config()

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

    UpperCAmelCase__	:Union[str, Any]      =      2
    UpperCAmelCase__	:List[str]      =      max_length

    for model_class in self.all_generative_model_classes:
       UpperCAmelCase__	:Union[str, Any]      =      model_class(__lowerCamelCase )

       UpperCAmelCase__	:Any      =      model.generate(__lowerCamelCase			,       attention_mask=__lowerCamelCase ).sequences
       self.assertEqual(generation_outputs.shape[-1]			,       __lowerCamelCase )

       UpperCAmelCase__	:List[Any]      =      jit(model.generate )
       UpperCAmelCase__	:str      =      jit_generate(__lowerCamelCase			,       attention_mask=__lowerCamelCase ).sequences

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




@require_flax
class 				UpperCAmelCase			(			unittest.TestCase						):


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[Any] ):
    UpperCAmelCase__	:List[Any]      =      AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-bert''' )
    UpperCAmelCase__	:Dict      =      FlaxAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-bert-flax-only''' )

    UpperCAmelCase__	:Union[str, Any]      =      '''Hello world'''
    UpperCAmelCase__	:List[str]      =      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''' ):
       UpperCAmelCase__	:Optional[int]      =      {'''foo''': '''bar'''}
       model.generate(__lowerCamelCase			,       **__lowerCamelCase )







'''simple docstring'''





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

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


if TYPE_CHECKING:
    from ... import FeatureExtractionMixin, TensorType

__lowerCamelCase												=						logging.get_logger(__name__)

__lowerCamelCase												=						{
    '''openai/imagegpt-small''': '''''',
    '''openai/imagegpt-medium''': '''''',
    '''openai/imagegpt-large''': '''''',
}




class 				UpperCAmelCase			(			_snake_case						):
 UpperCAmelCase				=					"imagegpt"
 UpperCAmelCase				=					["past_key_values"]
 UpperCAmelCase				=					{
     "hidden_size": "n_embd",
     "max_position_embeddings": "n_positions",
     "num_attention_heads": "n_head",
     "num_hidden_layers": "n_layer",
 }


 def __init__(      self							:	List[Any]			,       __lowerCamelCase							:	List[Any]=5_1_2 + 1			,       __lowerCamelCase							:	Dict=3_2 * 3_2			,       __lowerCamelCase							:	List[str]=5_1_2			,       __lowerCamelCase							:	List[Any]=2_4			,       __lowerCamelCase							:	Any=8			,       __lowerCamelCase							:	Tuple=None			,       __lowerCamelCase							:	Any="quick_gelu"			,       __lowerCamelCase							:	Optional[int]=0.1			,       __lowerCamelCase							:	str=0.1			,       __lowerCamelCase							:	int=0.1			,       __lowerCamelCase							:	Tuple=1e-5			,       __lowerCamelCase							:	Dict=0.02			,       __lowerCamelCase							:	Tuple=True			,       __lowerCamelCase							:	str=True			,       __lowerCamelCase							:	Union[str, Any]=False			,       __lowerCamelCase							:	Union[str, Any]=False			,       __lowerCamelCase							:	Dict=False			,       **__lowerCamelCase							:	str			,       ):
    UpperCAmelCase__	:Dict      =      vocab_size
    UpperCAmelCase__	:str      =      n_positions
    UpperCAmelCase__	:Tuple      =      n_embd
    UpperCAmelCase__	:Dict      =      n_layer
    UpperCAmelCase__	:List[Any]      =      n_head
    UpperCAmelCase__	:str      =      n_inner
    UpperCAmelCase__	:Optional[Any]      =      activation_function
    UpperCAmelCase__	:str      =      resid_pdrop
    UpperCAmelCase__	:Optional[Any]      =      embd_pdrop
    UpperCAmelCase__	:Tuple      =      attn_pdrop
    UpperCAmelCase__	:int      =      layer_norm_epsilon
    UpperCAmelCase__	:List[Any]      =      initializer_range
    UpperCAmelCase__	:List[Any]      =      scale_attn_weights
    UpperCAmelCase__	:List[str]      =      use_cache
    UpperCAmelCase__	:Tuple      =      scale_attn_by_inverse_layer_idx
    UpperCAmelCase__	:Union[str, Any]      =      reorder_and_upcast_attn
    UpperCAmelCase__	:List[Any]      =      tie_word_embeddings

    super().__init__(tie_word_embeddings=__lowerCamelCase			,       **__lowerCamelCase )




class 				UpperCAmelCase			(			_snake_case						):


 @property
 def        __SCREAMING_SNAKE_CASE					(      self							:	str ):
    return OrderedDict(
        [
            ('''input_ids''', {0: '''batch''', 1: '''sequence'''}),
        ] )





 def        __SCREAMING_SNAKE_CASE					(      self							:	Dict			,       __lowerCamelCase							:	"FeatureExtractionMixin"			,       __lowerCamelCase							:	int = 1			,       __lowerCamelCase							:	int = -1			,       __lowerCamelCase							:	bool = False			,       __lowerCamelCase							:	Optional["TensorType"] = None			,       __lowerCamelCase							:	int = 3			,       __lowerCamelCase							:	int = 3_2			,       __lowerCamelCase							:	int = 3_2			,       ):
    UpperCAmelCase__	:Tuple      =      self._generate_dummy_images(__lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase )
    UpperCAmelCase__	:Dict      =      dict(preprocessor(images=__lowerCamelCase			,       return_tensors=__lowerCamelCase ) )

    return inputs



