






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


  """simple docstring"""


  def __init__( self      ,  UpperCamelCase__      ,  UpperCamelCase__=14      ,  UpperCamelCase__=7      ,  UpperCamelCase__=True      ,  UpperCamelCase__=True      ,  UpperCamelCase__=False      ,  UpperCamelCase__=True      ,  UpperCamelCase__=99      ,  UpperCamelCase__=32      ,  UpperCamelCase__=5      ,  UpperCamelCase__=4      ,  UpperCamelCase__=37      ,  UpperCamelCase__="gelu"      ,  UpperCamelCase__=0.1      ,  UpperCamelCase__=0.1      ,  UpperCamelCase__=512      ,  UpperCamelCase__=0.02      ,  ):






         """simple docstring"""


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







  def       _a					( self  ):






         """simple docstring"""


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

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

         a_	    = 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=UpperCamelCase__      ,  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       _a					( self  ):






         """simple docstring"""


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







  def       _a					( self  ):






         """simple docstring"""


         a_    ,  a_    ,  a_	    = self.prepare_config_and_inputs()

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

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







  def       _a					( self      ,  UpperCamelCase__      ,  UpperCamelCase__      ,  UpperCamelCase__      ,  UpperCamelCase__  ):






         """simple docstring"""


         a_	    = 20
         a_	    = model_class_name(UpperCamelCase__  )

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

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

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

         a_	    = model(UpperCamelCase__  )

         a_	    = 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       _a					( self      ,  UpperCamelCase__      ,  UpperCamelCase__      ,  UpperCamelCase__      ,  UpperCamelCase__  ):






         """simple docstring"""


         a_	    = 20
         a_	    = model_class_name(UpperCamelCase__  )

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

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

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

         a_	    = model(UpperCamelCase__      ,  attention_mask=UpperCamelCase__  )

         a_	    = 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					):


  """simple docstring"""


  _a     :   Union[str, Any]			 =    (FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
  _a     :   Optional[Any]			 =    (FlaxGPTaLMHeadModel,) if is_flax_available() else ()







  def       _a					( self  ):






         """simple docstring"""


         a_	    = FlaxGPTaModelTester(self  )







  def       _a					( self  ):






         """simple docstring"""


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







  def       _a					( self  ):






         """simple docstring"""


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







  @slow
  def       _a					( self  ):






         """simple docstring"""


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

         a_	    = FlaxGPTaLMHeadModel.from_pretrained('gpt2'  )
         a_	    = False
         a_	    = model.config.eos_token_id

         a_	    = jax.jit(model.generate  )

         a_	    = jit_generate(inputs['input_ids']      ,  attention_mask=inputs['attention_mask']  ).sequences

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

         a_	    = [
             '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(UpperCamelCase__      ,  UpperCamelCase__  )







  @is_pt_flax_cross_test
  def       _a					( self  ):






         """simple docstring"""


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

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

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

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

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

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

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

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







  @is_pt_flax_cross_test
  def       _a					( self  ):






         """simple docstring"""


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

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

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

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

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

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

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

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

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

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







  @slow
  def       _a					( self  ):






         """simple docstring"""


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












'''simple docstring'''







def 			__UpperCamelCase				(			lowercase_					:     int      ,					lowercase_					:     int      ):
       """simple docstring"""




       return (input_a, input_a).count(0      ) % 2


def 			__UpperCamelCase				(			):
       """simple docstring"""




       assert xor_gate(0      ,					0      ) == 0
       assert xor_gate(0      ,					1      ) == 1
       assert xor_gate(1      ,					0      ) == 1
       assert xor_gate(1      ,					1      ) == 0


if __name__ == "__main__":
  print(xor_gate(0, 0))
  print(xor_gate(0, 1))




