





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 				_UpperCamelCase       :

       def __init__(		self:			List[Any]       , _SCREAMING_SNAKE_CASE:			Any       , _SCREAMING_SNAKE_CASE:			Union[str, Any]=14       , _SCREAMING_SNAKE_CASE:			Optional[int]=7       , _SCREAMING_SNAKE_CASE:			Optional[Any]=True       , _SCREAMING_SNAKE_CASE:			List[Any]=True       , _SCREAMING_SNAKE_CASE:			Optional[Any]=False       , _SCREAMING_SNAKE_CASE:			str=True       , _SCREAMING_SNAKE_CASE:			Optional[int]=99       , _SCREAMING_SNAKE_CASE:			Tuple=32       , _SCREAMING_SNAKE_CASE:			Optional[Any]=5       , _SCREAMING_SNAKE_CASE:			Tuple=4       , _SCREAMING_SNAKE_CASE:			int=37       , _SCREAMING_SNAKE_CASE:			Optional[Any]="gelu"       , _SCREAMING_SNAKE_CASE:			Any=0.1       , _SCREAMING_SNAKE_CASE:			List[str]=0.1       , _SCREAMING_SNAKE_CASE:			List[Any]=512       , _SCREAMING_SNAKE_CASE:			Optional[Any]=0.02       , )							->			Dict:




              """simple docstring"""


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

       def lowercase			(		self:			Dict		)							->			List[Any]:




              """simple docstring"""


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

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

              UpperCamelCase_ 							=	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=a_       , 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 lowercase			(		self:			List[Any]		)							->			str:




              """simple docstring"""


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

       def lowercase			(		self:			Optional[Any]		)							->			Dict:




              """simple docstring"""


              UpperCamelCase_							,			UpperCamelCase_							,			UpperCamelCase_ 							=	self.prepare_config_and_inputs()

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

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

       def lowercase			(		self:			str       , _SCREAMING_SNAKE_CASE:			Any       , _SCREAMING_SNAKE_CASE:			List[str]       , _SCREAMING_SNAKE_CASE:			Dict       , _SCREAMING_SNAKE_CASE:			List[str]		)							->			Tuple:




              """simple docstring"""


              UpperCamelCase_ 							=	20
              UpperCamelCase_ 							=	model_class_name(a_		)

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

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

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

              UpperCamelCase_ 							=	model(a_		)

              UpperCamelCase_ 							=	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 lowercase			(		self:			Any       , _SCREAMING_SNAKE_CASE:			Optional[int]       , _SCREAMING_SNAKE_CASE:			Optional[int]       , _SCREAMING_SNAKE_CASE:			Dict       , _SCREAMING_SNAKE_CASE:			Tuple		)							->			Any:




              """simple docstring"""


              UpperCamelCase_ 							=	20
              UpperCamelCase_ 							=	model_class_name(a_		)

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

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

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

              UpperCamelCase_ 							=	model(a_       , attention_mask=a_		)

              UpperCamelCase_ 							=	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 				_UpperCamelCase       (						_UpperCAmelCase	,    _UpperCAmelCase	,    unittest.TestCase  ):
       _UpperCamelCase					:       Optional[int]       = (FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
       _UpperCamelCase					:       Optional[Any]       = (FlaxGPTaLMHeadModel,) if is_flax_available() else ()

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




              """simple docstring"""


              UpperCamelCase_ 							=	FlaxGPTaModelTester(self		)

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




              """simple docstring"""


              for model_class_name in self.all_model_classes:
                     UpperCamelCase_							,			UpperCamelCase_							,			UpperCamelCase_ 							=	self.model_tester.prepare_config_and_inputs()
                     self.model_tester.check_use_cache_forward(a_       , a_       , a_       , a_		)

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




              """simple docstring"""


              for model_class_name in self.all_model_classes:
                     UpperCamelCase_							,			UpperCamelCase_							,			UpperCamelCase_ 							=	self.model_tester.prepare_config_and_inputs()
                     self.model_tester.check_use_cache_forward_with_attn_mask(
                         a_       , a_       , a_       , a_		)

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




              """simple docstring"""


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

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

              UpperCamelCase_ 							=	jax.jit(model.generate		)

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

              UpperCamelCase_ 							=	tokenizer.batch_decode(a_       , skip_special_tokens=a_		)

              UpperCamelCase_ 							=	[
                  "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(a_       , a_		)

       @is_pt_flax_cross_test
       def lowercase			(		self:			int		)							->			int:




              """simple docstring"""


              UpperCamelCase_							,			UpperCamelCase_ 							=	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
                            UpperCamelCase_ 							=	self._prepare_for_class(a_       , a_		)
                            UpperCamelCase_ 							=	{k: torch.tensor(v.tolist()		) for k, v in prepared_inputs_dict.items()}

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

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

                            UpperCamelCase_ 							=	convert_pytorch_state_dict_to_flax(pt_model.state_dict()       , a_		)
                            UpperCamelCase_ 							=	fx_state

                            with torch.no_grad():
                                   UpperCamelCase_ 							=	pt_model(**a_		).to_tuple()

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

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

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

       @is_pt_flax_cross_test
       def lowercase			(		self:			Dict		)							->			Union[str, Any]:




              """simple docstring"""


              UpperCamelCase_							,			UpperCamelCase_ 							=	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
                            UpperCamelCase_ 							=	self._prepare_for_class(a_       , a_		)
                            UpperCamelCase_ 							=	{k: torch.tensor(v.tolist()		) for k, v in prepared_inputs_dict.items()}

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

                            UpperCamelCase_ 							=	pt_model_class(a_		).eval()
                            UpperCamelCase_ 							=	model_class(a_       , dtype=jnp.floataa		)

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

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

                            with torch.no_grad():
                                   UpperCamelCase_ 							=	pt_model(**a_		).to_tuple()

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

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

                            with torch.no_grad():
                                   UpperCamelCase_ 							=	pt_model_loaded(**a_		).to_tuple()

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



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




              """simple docstring"""


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









"""simple docstring"""




import math






def 							lowerCAmelCase			(   UpperCamelCase_:				int       )							->				list:







  '''simple docstring'''







  if n <= 0:
    raise ValueError("Only positive integers have prime factors"       )
  _a					      =						[]
  while n % 2 == 0:
    pf.append(2       )
    _a					      =						int(n / 2       )
  for i in range(3  ,			int(math.sqrt(UpperCamelCase_       )       ) + 1  ,			2       ):
    while n % i == 0:
      pf.append(UpperCamelCase_       )
      _a					      =						int(n / i       )
  if n > 2:
    pf.append(UpperCamelCase_       )
  return pf






def 							lowerCAmelCase			(   UpperCamelCase_:				int       )							->				int:







  '''simple docstring'''







  if n <= 0:
    raise ValueError("Only positive numbers are accepted"       )
  _a					      =						1
  _a					      =						1
  while n % 2 == 0:
    temp += 1
    _a					      =						int(n / 2       )
  div *= temp
  for i in range(3  ,			int(math.sqrt(UpperCamelCase_       )       ) + 1  ,			2       ):
    _a					      =						1
    while n % i == 0:
      temp += 1
      _a					      =						int(n / i       )
    div *= temp
  if n > 1:
    div *= 2
  return div






def 							lowerCAmelCase			(   UpperCamelCase_:				int       )							->				int:







  '''simple docstring'''







  if n <= 0:
    raise ValueError("Only positive numbers are accepted"       )
  _a					      =						1
  _a					      =						1
  while n % 2 == 0:
    temp += 1
    _a					      =						int(n / 2       )
  if temp > 1:
    s *= (2**temp - 1) / (2 - 1)
  for i in range(3  ,			int(math.sqrt(UpperCamelCase_       )       ) + 1  ,			2       ):
    _a					      =						1
    while n % i == 0:
      temp += 1
      _a					      =						int(n / i       )
    if temp > 1:
      s *= (i**temp - 1) / (i - 1)
  return int(UpperCamelCase_       )






def 							lowerCAmelCase			(   UpperCamelCase_:				int       )							->				int:







  '''simple docstring'''







  _a					      =						n
  for x in set(prime_factors(UpperCamelCase_       )       ):
    s *= (x - 1) / x
  return int(UpperCamelCase_       )


if __name__ == "__main__":
     import doctest

     doctest.testmod()


