






import unittest

from parameterized import parameterized

from transformers import AutoTokenizer, GPTNeoXConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         GPTNeoXForCausalLM,
         GPTNeoXForQuestionAnswering,
         GPTNeoXForSequenceClassification,
         GPTNeoXForTokenClassification,
         GPTNeoXModel,
     )






class 						lowercase_   :






      '''simple docstring'''




      def __init__( self :		List[Any]					,   __UpperCAmelCase :		str					,   __UpperCAmelCase :		Optional[int]=13					,   __UpperCAmelCase :		List[str]=7					,   __UpperCAmelCase :		List[str]=True					,   __UpperCAmelCase :		List[str]=True					,   __UpperCAmelCase :		Optional[int]=True					,   __UpperCAmelCase :		int=True					,   __UpperCAmelCase :		List[str]=99					,   __UpperCAmelCase :		List[str]=64					,   __UpperCAmelCase :		Any=5					,   __UpperCAmelCase :		Any=4					,   __UpperCAmelCase :		Any=37					,   __UpperCAmelCase :		int="gelu"					,   __UpperCAmelCase :		Dict=0.1					,   __UpperCAmelCase :		int=0.1					,   __UpperCAmelCase :		Dict=512					,   __UpperCAmelCase :		List[str]=16					,   __UpperCAmelCase :		Optional[Any]=2					,   __UpperCAmelCase :		Tuple=0.02					,   __UpperCAmelCase :		Union[str, Any]=3					,   __UpperCAmelCase :		Union[str, Any]=4					,   __UpperCAmelCase :		List[str]=None					,   )       ->Union[str, Any]:




       """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						    = type_vocab_size
       a						    = type_sequence_label_size
       a						    = initializer_range
       a						    = num_labels
       a						    = num_choices
       a						    = scope
       a						    = vocab_size - 1







      def  __lowerCAmelCase   ( self :		Dict    )       ->int:




       """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						    = None
       if self.use_labels:
        a						    = ids_tensor([self.batch_size, self.seq_length]					,   self.num_labels    )

       a						    = self.get_config()

       return config, input_ids, input_mask, token_labels







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




       """simple docstring"""


       return GPTNeoXConfig(
           vocab_size=self.vocab_size					,   hidden_size=self.hidden_size					,   num_hidden_layers=self.num_hidden_layers					,   num_attention_heads=self.num_attention_heads					,   intermediate_size=self.intermediate_size					,   hidden_act=self.hidden_act					,   hidden_dropout_prob=self.hidden_dropout_prob					,   attention_probs_dropout_prob=self.attention_probs_dropout_prob					,   max_position_embeddings=self.max_position_embeddings					,   type_vocab_size=self.type_vocab_size					,   is_decoder=__UpperCAmelCase					,   initializer_range=self.initializer_range					,   pad_token_id=self.pad_token_id					,   )







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




       """simple docstring"""


       a    ,	a    ,	a    ,	a						    = self.prepare_config_and_inputs()

       a						    = True

       return config, input_ids, input_mask, token_labels







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




       """simple docstring"""


       a						    = GPTNeoXModel(config=__UpperCAmelCase    )
       model.to(__UpperCAmelCase    )
       model.eval()
       a						    = model(__UpperCAmelCase					,   attention_mask=__UpperCAmelCase    )
       a						    = model(__UpperCAmelCase    )
       self.parent.assertEqual(result.last_hidden_state.shape					,   (self.batch_size, self.seq_length, self.hidden_size)    )







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




       """simple docstring"""


       a						    = True
       a						    = GPTNeoXModel(__UpperCAmelCase    )
       model.to(__UpperCAmelCase    )
       model.eval()
       a						    = model(__UpperCAmelCase					,   attention_mask=__UpperCAmelCase    )
       self.parent.assertEqual(result.last_hidden_state.shape					,   (self.batch_size, self.seq_length, self.hidden_size)    )







      def  __lowerCAmelCase   ( self :		int					,   __UpperCAmelCase :		int					,   __UpperCAmelCase :		List[str]					,   __UpperCAmelCase :		Any					,   __UpperCAmelCase :		List[str]    )       ->Dict:




       """simple docstring"""


       a						    = GPTNeoXForCausalLM(config=__UpperCAmelCase    )
       model.to(__UpperCAmelCase    )
       model.eval()
       a						    = model(__UpperCAmelCase					,   attention_mask=__UpperCAmelCase					,   labels=__UpperCAmelCase    )
       self.parent.assertEqual(result.logits.shape					,   (self.batch_size, self.seq_length, self.vocab_size)    )







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




       """simple docstring"""


       a						    = self.num_labels
       a						    = GPTNeoXForQuestionAnswering(__UpperCAmelCase    )
       model.to(__UpperCAmelCase    )
       model.eval()
       a						    = model(__UpperCAmelCase					,   attention_mask=__UpperCAmelCase    )
       self.parent.assertEqual(result.start_logits.shape					,   (self.batch_size, self.seq_length)    )
       self.parent.assertEqual(result.end_logits.shape					,   (self.batch_size, self.seq_length)    )







      def  __lowerCAmelCase   ( self :		Dict					,   __UpperCAmelCase :		Tuple					,   __UpperCAmelCase :		Optional[Any]					,   __UpperCAmelCase :		str					,   __UpperCAmelCase :		Dict    )       ->Dict:




       """simple docstring"""


       a						    = self.num_labels
       a						    = GPTNeoXForSequenceClassification(__UpperCAmelCase    )
       model.to(__UpperCAmelCase    )
       model.eval()
       a						    = ids_tensor([self.batch_size]					,   self.type_sequence_label_size    )
       a						    = model(__UpperCAmelCase					,   attention_mask=__UpperCAmelCase					,   labels=__UpperCAmelCase    )
       self.parent.assertEqual(result.logits.shape					,   (self.batch_size, self.num_labels)    )







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




       """simple docstring"""


       a						    = self.num_labels
       a						    = GPTNeoXForTokenClassification(__UpperCAmelCase    )
       model.to(__UpperCAmelCase    )
       model.eval()
       a						    = model(__UpperCAmelCase					,   attention_mask=__UpperCAmelCase					,   labels=__UpperCAmelCase    )
       self.parent.assertEqual(result.logits.shape					,   (self.batch_size, self.seq_length, self.num_labels)    )







      def  __lowerCAmelCase   ( self :		str					,   __UpperCAmelCase :		Optional[Any]					,   __UpperCAmelCase :		Optional[int]					,   __UpperCAmelCase :		List[str]    )       ->str:




       """simple docstring"""


       a						    = True
       a						    = GPTNeoXForCausalLM(config=__UpperCAmelCase    )
       model.to(__UpperCAmelCase    )
       model.eval()

       # first forward pass
       a						    = model(__UpperCAmelCase					,   attention_mask=__UpperCAmelCase					,   use_cache=__UpperCAmelCase    )
       a						    = outputs.past_key_values

       # create hypothetical multiple next token and extent to next_input_ids
       a						    = ids_tensor((self.batch_size, 3)					,   config.vocab_size    )
       a						    = ids_tensor((self.batch_size, 3)					,   vocab_size=2    )

       # append to next input_ids and
       a						    = torch.cat([input_ids, next_tokens]					,   dim=-1    )
       a						    = torch.cat([input_mask, next_mask]					,   dim=-1    )

       a						    = model(__UpperCAmelCase					,   attention_mask=__UpperCAmelCase					,   output_hidden_states=__UpperCAmelCase    )
       a						    = output_from_no_past['''hidden_states'''][0]
       a						    = model(
           __UpperCAmelCase					,   attention_mask=__UpperCAmelCase					,   past_key_values=__UpperCAmelCase					,   output_hidden_states=__UpperCAmelCase					,   )['''hidden_states'''][0]

       # select random slice
       a						    = ids_tensor((1,)					,   output_from_past.shape[-1]    ).item()
       a						    = output_from_no_past[:, -3:, random_slice_idx].detach()
       a						    = output_from_past[:, :, random_slice_idx].detach()

       self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]    )

       # test that outputs are equal for slice
       self.parent.assertTrue(torch.allclose(__UpperCAmelCase					,   __UpperCAmelCase					,   atol=1e-3    )    )







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




       """simple docstring"""


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






@require_torch
class 						lowercase_   (      SCREAMING_SNAKE_CASE__					,    SCREAMING_SNAKE_CASE__					,    SCREAMING_SNAKE_CASE__					,    unittest.TestCase  ):






      '''simple docstring'''




      __snake_case	    =    (
          (
              GPTNeoXModel,
              GPTNeoXForCausalLM,
              GPTNeoXForQuestionAnswering,
              GPTNeoXForSequenceClassification,
              GPTNeoXForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      __snake_case	    =    (GPTNeoXForCausalLM,) if is_torch_available() else ()
      __snake_case	    =    (
          {
              '''feature-extraction''': GPTNeoXModel,
              '''question-answering''': GPTNeoXForQuestionAnswering,
              '''text-classification''': GPTNeoXForSequenceClassification,
              '''text-generation''': GPTNeoXForCausalLM,
              '''token-classification''': GPTNeoXForTokenClassification,
              '''zero-shot''': GPTNeoXForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      __snake_case	    =    False
      __snake_case	    =    False
      __snake_case	    =    False
      __snake_case	    =    False







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




       """simple docstring"""


       a						    = GPTNeoXModelTester(self    )
       a						    = ConfigTester(self					,   config_class=__UpperCAmelCase					,   hidden_size=64					,   num_attention_heads=8    )







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




       """simple docstring"""


       self.config_tester.run_common_tests()







      def  __lowerCAmelCase   ( self :		Dict    )       ->Dict:




       """simple docstring"""


       a    ,	a    ,	a    ,	a						    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(__UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase    )







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




       """simple docstring"""


       a    ,	a    ,	a    ,	a						    = self.model_tester.prepare_config_and_inputs_for_decoder()
       self.model_tester.create_and_check_model_as_decoder(__UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase    )







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




       """simple docstring"""


       a    ,	a    ,	a    ,	a						    = self.model_tester.prepare_config_and_inputs_for_decoder()

       a						    = None

       self.model_tester.create_and_check_model_as_decoder(__UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase    )







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




       """simple docstring"""


       a    ,	a    ,	a    ,	a						    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_decoder_model_past_large_inputs(__UpperCAmelCase					,   __UpperCAmelCase					,   __UpperCAmelCase    )







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




       """simple docstring"""


       a						    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_causal_lm(*__UpperCAmelCase    )







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




       """simple docstring"""


       a						    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*__UpperCAmelCase    )







      def  __lowerCAmelCase   ( self :		Tuple    )       ->Any:




       """simple docstring"""


       a						    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(*__UpperCAmelCase    )







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




       """simple docstring"""


       a						    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*__UpperCAmelCase    )







      @unittest.skip(reason='''Feed forward chunking is not implemented'''    )
      def  __lowerCAmelCase   ( self :		Union[str, Any]    )       ->Tuple:




       """simple docstring"""


       pass







      @parameterized.expand([('''linear''',), ('''dynamic''',)]    )
      def  __lowerCAmelCase   ( self :		str					,   __UpperCAmelCase :		str    )       ->Dict:




       """simple docstring"""


       a    ,	a						    = self.model_tester.prepare_config_and_inputs_for_common()
       a						    = ids_tensor([1, 10]					,   config.vocab_size    )
       a						    = ids_tensor([1, int(config.max_position_embeddings * 1.5    )]					,   config.vocab_size    )

       set_seed(42    )  # Fixed seed at init time so the two models get the same random weights
       a						    = GPTNeoXModel(__UpperCAmelCase    )
       original_model.to(__UpperCAmelCase    )
       original_model.eval()
       a						    = original_model(__UpperCAmelCase    ).last_hidden_state
       a						    = original_model(__UpperCAmelCase    ).last_hidden_state

       set_seed(42    )  # Fixed seed at init time so the two models get the same random weights
       a						    = {'''type''': scaling_type, '''factor''': 10.0}
       a						    = GPTNeoXModel(__UpperCAmelCase    )
       scaled_model.to(__UpperCAmelCase    )
       scaled_model.eval()
       a						    = scaled_model(__UpperCAmelCase    ).last_hidden_state
       a						    = scaled_model(__UpperCAmelCase    ).last_hidden_state

       # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
       # maximum sequence length, so the outputs for the short input should match.
       if scaling_type == "dynamic":
        self.assertTrue(torch.allclose(__UpperCAmelCase					,   __UpperCAmelCase					,   atol=1e-5    )    )
       else:
        self.assertFalse(torch.allclose(__UpperCAmelCase					,   __UpperCAmelCase					,   atol=1e-5    )    )

       # The output should be different for long inputs
       self.assertFalse(torch.allclose(__UpperCAmelCase					,   __UpperCAmelCase					,   atol=1e-5    )    )






@require_torch
class 						lowercase_   (      unittest.TestCase  ):






      '''simple docstring'''




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




       """simple docstring"""


       a						    = AutoTokenizer.from_pretrained('''EleutherAI/pythia-410m-deduped'''    )
       for checkpointing in [True, False]:
        a						    = GPTNeoXForCausalLM.from_pretrained('''EleutherAI/pythia-410m-deduped'''    )

        if checkpointing:
         model.gradient_checkpointing_enable()
        else:
         model.gradient_checkpointing_disable()
        model.to(__UpperCAmelCase    )

        a						    = tokenizer('''My favorite food is'''					,   return_tensors='''pt'''    ).to(__UpperCAmelCase    )
        # The hub repo. is updated on 2023-04-04, resulting in poor outputs.
        # See: https://github.com/huggingface/transformers/pull/24193
        a						    = '''My favorite food is a good old-fashioned, old-fashioned, old-fashioned.\n\nI\'m not sure'''

        a						    = model.generate(**__UpperCAmelCase					,   do_sample=__UpperCAmelCase					,   max_new_tokens=20    )
        a						    = tokenizer.batch_decode(__UpperCAmelCase    )[0]

        self.assertEqual(__UpperCAmelCase					,   __UpperCAmelCase    )








import unittest

from diffusers import FlaxAutoencoderKL
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax

from .test_modeling_common_flax import FlaxModelTesterMixin


if is_flax_available():
				import jax



@require_flax
class       a_       (     SCREAMING_SNAKE_CASE__  ,    unittest.TestCase     ):
	A    							= FlaxAutoencoderKL





	@property
	def 						A_(							self      )  ->					Optional[Any]:




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						4
				SCREAMING_SNAKE_CASE_								=						3
				SCREAMING_SNAKE_CASE_								=						(32, 32)

				SCREAMING_SNAKE_CASE_								=						jax.random.PRNGKey(0      )
				SCREAMING_SNAKE_CASE_								=						jax.random.uniform(SCREAMING_SNAKE_CASE      ,     ((batch_size, num_channels) + sizes)      )

				return {"sample": image, "prng_key": prng_key}





	def 						A_(							self      )  ->					Tuple:




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						{
				    'block_out_channels': [32, 64],
				    'in_channels': 3,
				    'out_channels': 3,
				    'down_block_types': ['DownEncoderBlock2D', 'DownEncoderBlock2D'],
				    'up_block_types': ['UpDecoderBlock2D', 'UpDecoderBlock2D'],
				    'latent_channels': 4,
				}
				SCREAMING_SNAKE_CASE_								=						self.dummy_input
				return init_dict, inputs_dict
