





import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTNeoConfig, 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, 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.gpt_neo.modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel

if is_torch_available():
  import torch



class 				__snake_case  :


    '''simple docstring'''







    def __init__(					self      ,							A_      ,							A_=14      ,							A_=7      ,							A_=True      ,							A_=True      ,							A_=False      ,							A_=True      ,							A_=99      ,							A_=32      ,							A_=4      ,							A_=4      ,							A_=[[["global", "local"], 2]]      ,							A_=37      ,							A_="gelu"      ,							A_=0.1      ,							A_=0.1      ,							A_=5_12      ,							A_=7      ,							A_=0.02      ,							):





        '''simple docstring'''

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



    def    lowercase_				(					self    ):





        '''simple docstring'''

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

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

        SCREAMING_SNAKE_CASE__       							= GPTNeoConfig(
            vocab_size=self.vocab_size      ,							hidden_size=self.hidden_size      ,							num_layers=self.num_hidden_layers      ,							num_heads=self.num_attention_heads      ,							max_position_embeddings=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      ,							window_size=self.window_size      ,							attention_types=self.attention_types      ,							)

        return (config, input_ids, input_mask)



    def    lowercase_				(					self    ):





        '''simple docstring'''

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



    def    lowercase_				(					self      ,							A_      ,							A_      ,							A_      ,							A_    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= 20
        SCREAMING_SNAKE_CASE__       							= model_class_name(A_    )

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

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

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

        SCREAMING_SNAKE_CASE__       							= model(A_    )

        SCREAMING_SNAKE_CASE__       							= 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      ,							A_      ,							A_      ,							A_      ,							A_    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= 20
        SCREAMING_SNAKE_CASE__       							= model_class_name(A_    )

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

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

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

        SCREAMING_SNAKE_CASE__       							= model(A_      ,							attention_mask=A_    )

        SCREAMING_SNAKE_CASE__       							= 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 				__snake_case  (     __SCREAMING_SNAKE_CASE ,    __SCREAMING_SNAKE_CASE ,    unittest.TestCase						):


    '''simple docstring'''







    lowerCamelCase__						:		Any							=  (FlaxGPTNeoModel, FlaxGPTNeoForCausalLM) if is_flax_available() else ()
    lowerCamelCase__						:		int							=  (FlaxGPTNeoForCausalLM,) if is_flax_available() else ()



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= FlaxGPTNeoModelTester(self    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        for model_class_name in self.all_model_classes:
            SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__       							= self.model_tester.prepare_config_and_inputs()
            self.model_tester.check_use_cache_forward(A_      ,							A_      ,							A_      ,							A_    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        for model_class_name in self.all_model_classes:
            SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__       							= 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    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= GPTaTokenizer.from_pretrained('''gpt2'''      ,							pad_token='''<|endoftext|>'''      ,							padding_side='''left'''    )
        SCREAMING_SNAKE_CASE__       							= tokenizer(['''Hello this is a long string''', '''Hey''']      ,							return_tensors='''np'''      ,							padding=A_      ,							truncation=A_    )

        SCREAMING_SNAKE_CASE__       							= FlaxGPTNeoForCausalLM.from_pretrained('''EleutherAI/gpt-neo-125M'''    )
        SCREAMING_SNAKE_CASE__       							= False
        SCREAMING_SNAKE_CASE__       							= model.config.eos_token_id

        SCREAMING_SNAKE_CASE__       							= jax.jit(model.generate    )

        SCREAMING_SNAKE_CASE__       							= jit_generate(
            inputs['''input_ids''']      ,							attention_mask=inputs['''attention_mask''']      ,							pad_token_id=tokenizer.pad_token_id    ).sequences

        SCREAMING_SNAKE_CASE__       							= tokenizer.batch_decode(A_      ,							skip_special_tokens=A_    )

        SCREAMING_SNAKE_CASE__       							= [
            '''Hello this is a long string of text.\n\nI\'m trying to get the text of the''',
            '''Hey, I\'m a little late to the party. I\'m going to''',
        ]

        self.assertListEqual(A_      ,							A_    )



    @is_pt_flax_cross_test
    def    lowercase_				(					self    ):





        '''simple docstring'''

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

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

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

                SCREAMING_SNAKE_CASE__       							= convert_pytorch_state_dict_to_flax(pt_model.state_dict()      ,							A_    )
                SCREAMING_SNAKE_CASE__       							= fx_state

                with torch.no_grad():
                    SCREAMING_SNAKE_CASE__       							= pt_model(**A_    ).to_tuple()

                SCREAMING_SNAKE_CASE__       							= 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_    )
                    SCREAMING_SNAKE_CASE__       							= model_class.from_pretrained(A_      ,							from_pt=A_    )

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





        '''simple docstring'''

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

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

                SCREAMING_SNAKE_CASE__       							= pt_model_class(A_    ).eval()
                SCREAMING_SNAKE_CASE__       							= model_class(A_      ,							dtype=jnp.floataa    )

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

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

                with torch.no_grad():
                    SCREAMING_SNAKE_CASE__       							= pt_model(**A_    ).to_tuple()

                SCREAMING_SNAKE_CASE__       							= 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_    )
                    SCREAMING_SNAKE_CASE__       							= pt_model_class.from_pretrained(A_      ,							from_flax=A_    )

                with torch.no_grad():
                    SCREAMING_SNAKE_CASE__       							= 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    ):





        '''simple docstring'''

        for model_class_name in self.all_model_classes:
            SCREAMING_SNAKE_CASE__       							= model_class_name.from_pretrained('''EleutherAI/gpt-neo-125M'''    )
            SCREAMING_SNAKE_CASE__       							= model(np.ones((1, 1)    )    )
            self.assertIsNotNone(A_    )







import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
  import transformer_engine.pytorch as te


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_=True   ,       lowerCAmelCase_=True   ,       lowerCAmelCase_=True  )							->	Tuple:
    if not is_fpa_available():
        raise ImportError('''Using `convert_model` requires transformer_engine to be installed.'''  )
    for name, module in model.named_children():
        if isinstance(lowerCAmelCase_   ,       nn.Linear  ) and to_transformer_engine and _convert_linear:
            # Return early if the linear layer weights are not multiples of 16
            if any(p % 1_6 != 0 for p in module.weight.shape  ):
                return
            SCREAMING_SNAKE_CASE__       							= module.bias is not None
            SCREAMING_SNAKE_CASE__       							= te.Linear(
                module.in_features   ,       module.out_features   ,       bias=lowerCAmelCase_   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            if has_bias:
                SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

            setattr(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        elif isinstance(lowerCAmelCase_   ,       nn.LayerNorm  ) and to_transformer_engine and _convert_ln:
            SCREAMING_SNAKE_CASE__       							= te.LayerNorm(module.normalized_shape[0]   ,       eps=module.eps   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

            setattr(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        elif isinstance(lowerCAmelCase_   ,       te.Linear  ) and not to_transformer_engine and _convert_linear:
            SCREAMING_SNAKE_CASE__       							= module.bias is not None
            SCREAMING_SNAKE_CASE__       							= nn.Linear(
                module.in_features   ,       module.out_features   ,       bias=lowerCAmelCase_   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            if has_bias:
                SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

            setattr(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        elif isinstance(lowerCAmelCase_   ,       te.LayerNorm  ) and not to_transformer_engine and _convert_ln:
            SCREAMING_SNAKE_CASE__       							= nn.LayerNorm(module.normalized_shape[0]   ,       eps=module.eps   ,       params_dtype=module.weight.dtype  )
            SCREAMING_SNAKE_CASE__       							= module.weight.data.clone()
            SCREAMING_SNAKE_CASE__       							= module.bias.data.clone()

            setattr(lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )
        else:
            convert_model(
                lowerCAmelCase_   ,       to_transformer_engine=lowerCAmelCase_   ,       _convert_linear=_convert_linear   ,       _convert_ln=_convert_ln   ,       )



def        __snake_case       (						lowerCAmelCase_  )							->	Union[str, Any]:
    if not is_fpa_available():
        raise ImportError('''Using `has_transformer_engine_layers` requires transformer_engine to be installed.'''  )
    for m in model.modules():
        if isinstance(lowerCAmelCase_   ,       (te.LayerNorm, te.Linear)  ):
            return True
    return False
