




'''simple docstring'''
import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaConfig, GPTaTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
  import jax
  import jax.numpy as jnp

  from transformers.modeling_flax_pytorch_utils import (
      convert_pytorch_state_dict_to_flax,
      load_flax_weights_in_pytorch_model,
  )
  from transformers.models.gpta.modeling_flax_gpta import FlaxGPTaLMHeadModel, FlaxGPTaModel

if is_torch_available():
  import torch






class        SCREAMING_SNAKE_CASE:

       def __init__(    self					,		lowerCamelCase__					,		lowerCamelCase__=14					,		lowerCamelCase__=7					,		lowerCamelCase__=True					,		lowerCamelCase__=True					,		lowerCamelCase__=False					,		lowerCamelCase__=True					,		lowerCamelCase__=99					,		lowerCamelCase__=32					,		lowerCamelCase__=5					,		lowerCamelCase__=4					,		lowerCamelCase__=37					,		lowerCamelCase__="gelu"					,		lowerCamelCase__=0.1					,		lowerCamelCase__=0.1					,		lowerCamelCase__=512					,		lowerCamelCase__=0.02					,		)    ->					List[Any]:

           """simple docstring"""



           __lowercase							=						parent
           __lowercase							=						batch_size
           __lowercase							=						seq_length
           __lowercase							=						is_training
           __lowercase							=						use_input_mask
           __lowercase							=						use_token_type_ids
           __lowercase							=						use_labels
           __lowercase							=						vocab_size
           __lowercase							=						hidden_size
           __lowercase							=						num_hidden_layers
           __lowercase							=						num_attention_heads
           __lowercase							=						intermediate_size
           __lowercase							=						hidden_act
           __lowercase							=						hidden_dropout_prob
           __lowercase							=						attention_probs_dropout_prob
           __lowercase							=						max_position_embeddings
           __lowercase							=						initializer_range
           __lowercase							=						None
           __lowercase							=						vocab_size - 1
           __lowercase							=						vocab_size - 1
           __lowercase							=						vocab_size - 1

       def     snake_case__			(    self     )    ->					List[Any]:

           """simple docstring"""



           __lowercase							=						ids_tensor([self.batch_size, self.seq_length]					,		self.vocab_size     )

           __lowercase							=						None
           if self.use_input_mask:
               __lowercase							=						random_attention_mask([self.batch_size, self.seq_length]     )

           __lowercase							=						GPTaConfig(
               vocab_size=self.vocab_size					,		n_embd=self.hidden_size					,		n_layer=self.num_hidden_layers					,		n_head=self.num_attention_heads					,		n_positions=self.max_position_embeddings					,		use_cache=lowerCamelCase__					,		bos_token_id=self.bos_token_id					,		eos_token_id=self.eos_token_id					,		pad_token_id=self.pad_token_id					,		)

           return (config, input_ids, input_mask)

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

           """simple docstring"""



           __lowercase							=						self.prepare_config_and_inputs()
           __lowercase				,__lowercase				,__lowercase							=						config_and_inputs
           __lowercase							=						{"""input_ids""": input_ids, """attention_mask""": attention_mask}
           return config, inputs_dict

       def     snake_case__			(    self     )    ->					int:

           """simple docstring"""



           __lowercase				,__lowercase				,__lowercase							=						self.prepare_config_and_inputs()

           __lowercase							=						floats_tensor([self.batch_size, self.seq_length, self.hidden_size]     )
           __lowercase							=						ids_tensor([self.batch_size, self.seq_length]					,		vocab_size=2     )

           return (
               config,
               input_ids,
               attention_mask,
               encoder_hidden_states,
               encoder_attention_mask,
           )

       def     snake_case__			(    self					,		lowerCamelCase__					,		lowerCamelCase__					,		lowerCamelCase__					,		lowerCamelCase__     )    ->					Optional[Any]:

           """simple docstring"""



           __lowercase							=						20
           __lowercase							=						model_class_name(lowerCamelCase__     )

           __lowercase							=						model.init_cache(input_ids.shape[0]					,		lowerCamelCase__     )
           __lowercase							=						jnp.ones((input_ids.shape[0], max_decoder_length)					,		dtype="""i4"""     )

           __lowercase							=						jnp.broadcast_to(
               jnp.arange(input_ids.shape[-1] - 1     )[None, :]					,		(input_ids.shape[0], input_ids.shape[-1] - 1)     )
           __lowercase							=						model(
               input_ids[:, :-1]					,		attention_mask=lowerCamelCase__					,		past_key_values=lowerCamelCase__					,		position_ids=lowerCamelCase__					,		)

           __lowercase							=						jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]					,		dtype="""i4"""     )
           __lowercase							=						model(
               input_ids[:, -1:]					,		attention_mask=lowerCamelCase__					,		past_key_values=outputs_cache.past_key_values					,		position_ids=lowerCamelCase__					,		)

           __lowercase							=						model(lowerCamelCase__     )

           __lowercase							=						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     snake_case__			(    self					,		lowerCamelCase__					,		lowerCamelCase__					,		lowerCamelCase__					,		lowerCamelCase__     )    ->					Union[str, Any]:

           """simple docstring"""



           __lowercase							=						20
           __lowercase							=						model_class_name(lowerCamelCase__     )

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

           __lowercase							=						model.init_cache(input_ids.shape[0]					,		lowerCamelCase__     )
           __lowercase							=						jnp.broadcast_to(
               jnp.arange(input_ids.shape[-1] - 1     )[None, :]					,		(input_ids.shape[0], input_ids.shape[-1] - 1)     )

           __lowercase							=						model(
               input_ids[:, :-1]					,		attention_mask=lowerCamelCase__					,		past_key_values=lowerCamelCase__					,		position_ids=lowerCamelCase__					,		)
           __lowercase							=						jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]					,		dtype="""i4"""     )
           __lowercase							=						model(
               input_ids[:, -1:]					,		past_key_values=outputs_cache.past_key_values					,		attention_mask=lowerCamelCase__					,		position_ids=lowerCamelCase__					,		)

           __lowercase							=						model(lowerCamelCase__					,		attention_mask=lowerCamelCase__     )

           __lowercase							=						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        SCREAMING_SNAKE_CASE(				__A   ,      __A   ,      unittest.TestCase     ):
       snake_case_ :     Any    =					(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
       snake_case_ :     List[Any]    =					(FlaxGPTaLMHeadModel,) if is_flax_available() else ()

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

           """simple docstring"""



           __lowercase							=						FlaxGPTaModelTester(self     )

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

           """simple docstring"""



           for model_class_name in self.all_model_classes:
               __lowercase				,__lowercase				,__lowercase							=						self.model_tester.prepare_config_and_inputs()
               self.model_tester.check_use_cache_forward(lowerCamelCase__					,		lowerCamelCase__					,		lowerCamelCase__					,		lowerCamelCase__     )

       def     snake_case__			(    self     )    ->					Dict:

           """simple docstring"""



           for model_class_name in self.all_model_classes:
               __lowercase				,__lowercase				,__lowercase							=						self.model_tester.prepare_config_and_inputs()
               self.model_tester.check_use_cache_forward_with_attn_mask(
                   lowerCamelCase__					,		lowerCamelCase__					,		lowerCamelCase__					,		lowerCamelCase__     )

       @slow
       def     snake_case__			(    self     )    ->					Optional[int]:

           """simple docstring"""



           __lowercase							=						GPTaTokenizer.from_pretrained("""gpt2"""					,		pad_token="""</s>"""					,		padding_side="""left"""     )
           __lowercase							=						tokenizer(["""Hello this is a long string""", """Hey"""]					,		return_tensors="""np"""					,		padding=lowerCamelCase__					,		truncation=lowerCamelCase__     )

           __lowercase							=						FlaxGPTaLMHeadModel.from_pretrained("""gpt2"""     )
           __lowercase							=						False
           __lowercase							=						model.config.eos_token_id

           __lowercase							=						jax.jit(model.generate     )

           __lowercase							=						jit_generate(inputs["""input_ids"""]					,		attention_mask=inputs["""attention_mask"""]     ).sequences

           __lowercase							=						tokenizer.batch_decode(lowerCamelCase__					,		skip_special_tokens=lowerCamelCase__     )

           __lowercase							=						[
               """Hello this is a long string of words. I'm going to start with the first one.\n""",
               """Hey, I'm not sure if I'm going to be able to do""",
           ]

           self.assertListEqual(lowerCamelCase__					,		lowerCamelCase__     )

       @is_pt_flax_cross_test
       def     snake_case__			(    self     )    ->					str:

           """simple docstring"""



           __lowercase				,__lowercase							=						self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               with self.subTest(model_class.__name__     ):
                   # prepare inputs
                   __lowercase							=						self._prepare_for_class(lowerCamelCase__					,		lowerCamelCase__     )
                   __lowercase							=						{k: torch.tensor(v.tolist()     ) for k, v in prepared_inputs_dict.items()}

                   # load corresponding PyTorch class
                   __lowercase							=						model_class.__name__[4:]  # Skip the "Flax" at the beginning
                   __lowercase							=						getattr(lowerCamelCase__					,		lowerCamelCase__     )

                   __lowercase				,__lowercase							=						pt_inputs["""input_ids"""].shape
                   __lowercase							=						np.random.randint(0					,		seq_length - 1					,		size=(batch_size,)     )
                   for batch_idx, start_index in enumerate(lowerCamelCase__     ):
                       __lowercase							=						0
                       __lowercase							=						1
                       __lowercase							=						0
                       __lowercase							=						1
                   __lowercase							=						pt_model_class(lowerCamelCase__     ).eval()
                   __lowercase							=						model_class(lowerCamelCase__					,		dtype=jnp.floataa     )

                   __lowercase							=						convert_pytorch_state_dict_to_flax(pt_model.state_dict()					,		lowerCamelCase__     )
                   __lowercase							=						fx_state

                   with torch.no_grad():
                       __lowercase							=						pt_model(**lowerCamelCase__     ).to_tuple()

                   __lowercase							=						fx_model(**lowerCamelCase__     ).to_tuple()
                   self.assertEqual(len(lowerCamelCase__     )					,		len(lowerCamelCase__     )					,		"""Output lengths differ between Flax and PyTorch"""     )
                   for fx_output, pt_output in zip(lowerCamelCase__					,		lowerCamelCase__     ):
                       self.assert_almost_equals(fx_output[:, -1]					,		pt_output[:, -1].numpy()					,		4E-2     )

                   with tempfile.TemporaryDirectory() as tmpdirname:
                       pt_model.save_pretrained(lowerCamelCase__     )
                       __lowercase							=						model_class.from_pretrained(lowerCamelCase__					,		from_pt=lowerCamelCase__     )

                   __lowercase							=						fx_model_loaded(**lowerCamelCase__     ).to_tuple()
                   self.assertEqual(
                       len(lowerCamelCase__     )					,		len(lowerCamelCase__     )					,		"""Output lengths differ between Flax and PyTorch"""     )
                   for fx_output_loaded, pt_output in zip(lowerCamelCase__					,		lowerCamelCase__     ):
                       self.assert_almost_equals(fx_output_loaded[:, -1]					,		pt_output[:, -1].numpy()					,		4E-2     )

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

           """simple docstring"""



           __lowercase				,__lowercase							=						self.model_tester.prepare_config_and_inputs_for_common()
           for model_class in self.all_model_classes:
               with self.subTest(model_class.__name__     ):
                   # prepare inputs
                   __lowercase							=						self._prepare_for_class(lowerCamelCase__					,		lowerCamelCase__     )
                   __lowercase							=						{k: torch.tensor(v.tolist()     ) for k, v in prepared_inputs_dict.items()}

                   # load corresponding PyTorch class
                   __lowercase							=						model_class.__name__[4:]  # Skip the "Flax" at the beginning
                   __lowercase							=						getattr(lowerCamelCase__					,		lowerCamelCase__     )

                   __lowercase							=						pt_model_class(lowerCamelCase__     ).eval()
                   __lowercase							=						model_class(lowerCamelCase__					,		dtype=jnp.floataa     )

                   __lowercase							=						load_flax_weights_in_pytorch_model(lowerCamelCase__					,		fx_model.params     )
                   __lowercase				,__lowercase							=						pt_inputs["""input_ids"""].shape
                   __lowercase							=						np.random.randint(0					,		seq_length - 1					,		size=(batch_size,)     )
                   for batch_idx, start_index in enumerate(lowerCamelCase__     ):
                       __lowercase							=						0
                       __lowercase							=						1
                       __lowercase							=						0
                       __lowercase							=						1

                   # make sure weights are tied in PyTorch
                   pt_model.tie_weights()

                   with torch.no_grad():
                       __lowercase							=						pt_model(**lowerCamelCase__     ).to_tuple()

                   __lowercase							=						fx_model(**lowerCamelCase__     ).to_tuple()
                   self.assertEqual(len(lowerCamelCase__     )					,		len(lowerCamelCase__     )					,		"""Output lengths differ between Flax and PyTorch"""     )
                   for fx_output, pt_output in zip(lowerCamelCase__					,		lowerCamelCase__     ):
                       self.assert_almost_equals(fx_output[:, -1]					,		pt_output[:, -1].numpy()					,		4E-2     )

                   with tempfile.TemporaryDirectory() as tmpdirname:
                       fx_model.save_pretrained(lowerCamelCase__     )
                       __lowercase							=						pt_model_class.from_pretrained(lowerCamelCase__					,		from_flax=lowerCamelCase__     )

                   with torch.no_grad():
                       __lowercase							=						pt_model_loaded(**lowerCamelCase__     ).to_tuple()

                   self.assertEqual(
                       len(lowerCamelCase__     )					,		len(lowerCamelCase__     )					,		"""Output lengths differ between Flax and PyTorch"""     )
                   for fx_output, pt_output in zip(lowerCamelCase__					,		lowerCamelCase__     ):
                       self.assert_almost_equals(fx_output[:, -1]					,		pt_output[:, -1].numpy()					,		4E-2     )




       @slow
       def     snake_case__			(    self     )    ->					List[Any]:

           """simple docstring"""



           for model_class_name in self.all_model_classes:
               __lowercase							=						model_class_name.from_pretrained("""gpt2"""					,		from_pt=lowerCamelCase__     )
               __lowercase							=						model(np.ones((1, 1)     )     )
               self.assertIsNotNone(lowerCamelCase__     )










'''simple docstring'''
def 				snake_case_  (   a__		:	int			):



    """simple docstring"""




    if not isinstance(a__    ,a__			) or number < 0:
        raise ValueError("""Input must be a non-negative integer"""			)

    __lowercase							=						number & ~(number - 1)
    __lowercase							=						0
    while intermediate:
        intermediate >>= 1
        index += 1
    return index - 1


if __name__ == "__main__":
  import doctest

  doctest.testmod(verbose=True)




