





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







  def __init__(						self     ,    _lowercase     ,    _lowercase=14     ,    _lowercase=7     ,    _lowercase=True     ,    _lowercase=True     ,    _lowercase=False     ,    _lowercase=True     ,    _lowercase=99     ,    _lowercase=32     ,    _lowercase=5     ,    _lowercase=4     ,    _lowercase=37     ,    _lowercase="gelu"     ,    _lowercase=0.1     ,    _lowercase=0.1     ,    _lowercase=512     ,    _lowercase=0.02     ,    ):
       lowerCAmelCase_       :							Optional[Any]			 =		parent
       lowerCAmelCase_       :							List[str]			 =		batch_size
       lowerCAmelCase_       :							Dict			 =		seq_length
       lowerCAmelCase_       :							Tuple			 =		is_training
       lowerCAmelCase_       :							int			 =		use_input_mask
       lowerCAmelCase_       :							Optional[Any]			 =		use_token_type_ids
       lowerCAmelCase_       :							List[Any]			 =		use_labels
       lowerCAmelCase_       :							Dict			 =		vocab_size
       lowerCAmelCase_       :							Optional[int]			 =		hidden_size
       lowerCAmelCase_       :							List[Any]			 =		num_hidden_layers
       lowerCAmelCase_       :							Any			 =		num_attention_heads
       lowerCAmelCase_       :							List[Any]			 =		intermediate_size
       lowerCAmelCase_       :							str			 =		hidden_act
       lowerCAmelCase_       :							int			 =		hidden_dropout_prob
       lowerCAmelCase_       :							str			 =		attention_probs_dropout_prob
       lowerCAmelCase_       :							int			 =		max_position_embeddings
       lowerCAmelCase_       :							List[Any]			 =		initializer_range
       lowerCAmelCase_       :							Union[str, Any]			 =		None
       lowerCAmelCase_       :							Optional[int]			 =		vocab_size - 1
       lowerCAmelCase_       :							Dict			 =		vocab_size - 1
       lowerCAmelCase_       :							Union[str, Any]			 =		vocab_size - 1







  def 						UpperCAmelCase__		(						self    ):
       lowerCAmelCase_       :							List[str]			 =		ids_tensor([self.batch_size, self.seq_length]     ,    self.vocab_size    )

       lowerCAmelCase_       :							Dict			 =		None
       if self.use_input_mask:
            lowerCAmelCase_       :							Any			 =		random_attention_mask([self.batch_size, self.seq_length]    )

       lowerCAmelCase_       :							Dict			 =		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=_lowercase     ,    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 						UpperCAmelCase__		(						self    ):
       lowerCAmelCase_       :							Optional[Any]			 =		self.prepare_config_and_inputs()
       lowerCAmelCase_       :							Any			 =		config_and_inputs
       lowerCAmelCase_       :							Dict			 =		{"input_ids": input_ids, "attention_mask": attention_mask}
       return config, inputs_dict







  def 						UpperCAmelCase__		(						self    ):
       lowerCAmelCase_       :							List[str]			 =		self.prepare_config_and_inputs()

       lowerCAmelCase_       :							Optional[int]			 =		floats_tensor([self.batch_size, self.seq_length, self.hidden_size]    )
       lowerCAmelCase_       :							Optional[int]			 =		ids_tensor([self.batch_size, self.seq_length]     ,    vocab_size=2    )

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







  def 						UpperCAmelCase__		(						self     ,    _lowercase     ,    _lowercase     ,    _lowercase     ,    _lowercase    ):
       lowerCAmelCase_       :							List[str]			 =		20
       lowerCAmelCase_       :							str			 =		model_class_name(_lowercase    )

       lowerCAmelCase_       :							Dict			 =		model.init_cache(input_ids.shape[0]     ,    _lowercase    )
       lowerCAmelCase_       :							List[Any]			 =		jnp.ones((input_ids.shape[0], max_decoder_length)     ,    dtype="""i4"""    )

       lowerCAmelCase_       :							Optional[Any]			 =		jnp.broadcast_to(
           jnp.arange(input_ids.shape[-1] - 1    )[None, :]     ,    (input_ids.shape[0], input_ids.shape[-1] - 1)    )
       lowerCAmelCase_       :							Union[str, Any]			 =		model(
           input_ids[:, :-1]     ,    attention_mask=_lowercase     ,    past_key_values=_lowercase     ,    position_ids=_lowercase     ,    )

       lowerCAmelCase_       :							Optional[int]			 =		jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]     ,    dtype="""i4"""    )
       lowerCAmelCase_       :							List[Any]			 =		model(
           input_ids[:, -1:]     ,    attention_mask=_lowercase     ,    past_key_values=outputs_cache.past_key_values     ,    position_ids=_lowercase     ,    )

       lowerCAmelCase_       :							Optional[Any]			 =		model(_lowercase    )

       lowerCAmelCase_       :							int			 =		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 						UpperCAmelCase__		(						self     ,    _lowercase     ,    _lowercase     ,    _lowercase     ,    _lowercase    ):
       lowerCAmelCase_       :							Optional[Any]			 =		20
       lowerCAmelCase_       :							Union[str, Any]			 =		model_class_name(_lowercase    )

       lowerCAmelCase_       :							Optional[int]			 =		jnp.concatenate(
           [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])    )]     ,    axis=-1     ,    )

       lowerCAmelCase_       :							List[str]			 =		model.init_cache(input_ids.shape[0]     ,    _lowercase    )
       lowerCAmelCase_       :							Optional[Any]			 =		jnp.broadcast_to(
           jnp.arange(input_ids.shape[-1] - 1    )[None, :]     ,    (input_ids.shape[0], input_ids.shape[-1] - 1)    )

       lowerCAmelCase_       :							Dict			 =		model(
           input_ids[:, :-1]     ,    attention_mask=_lowercase     ,    past_key_values=_lowercase     ,    position_ids=_lowercase     ,    )
       lowerCAmelCase_       :							Tuple			 =		jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]     ,    dtype="""i4"""    )
       lowerCAmelCase_       :							List[Any]			 =		model(
           input_ids[:, -1:]     ,    past_key_values=outputs_cache.past_key_values     ,    attention_mask=_lowercase     ,    position_ids=_lowercase     ,    )

       lowerCAmelCase_       :							List[str]			 =		model(_lowercase     ,    attention_mask=_lowercase    )

       lowerCAmelCase_       :							int			 =		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   lowercase__     (				__A  ,       __A  ,       unittest.TestCase	):
  __UpperCamelCase						  =	(FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
  __UpperCamelCase						  =	(FlaxGPTaLMHeadModel,) if is_flax_available() else ()







  def 						UpperCAmelCase__		(						self    ):
       lowerCAmelCase_       :							str			 =		FlaxGPTaModelTester(self    )







  def 						UpperCAmelCase__		(						self    ):
       for model_class_name in self.all_model_classes:
            lowerCAmelCase_       :							str			 =		self.model_tester.prepare_config_and_inputs()
            self.model_tester.check_use_cache_forward(_lowercase     ,    _lowercase     ,    _lowercase     ,    _lowercase    )







  def 						UpperCAmelCase__		(						self    ):
       for model_class_name in self.all_model_classes:
            lowerCAmelCase_       :							List[Any]			 =		self.model_tester.prepare_config_and_inputs()
            self.model_tester.check_use_cache_forward_with_attn_mask(
                _lowercase     ,    _lowercase     ,    _lowercase     ,    _lowercase    )







  @slow
  def 						UpperCAmelCase__		(						self    ):
       lowerCAmelCase_       :							List[Any]			 =		GPTaTokenizer.from_pretrained("""gpt2"""     ,    pad_token="""</s>"""     ,    padding_side="""left"""    )
       lowerCAmelCase_       :							List[Any]			 =		tokenizer(["""Hello this is a long string""", """Hey"""]     ,    return_tensors="""np"""     ,    padding=_lowercase     ,    truncation=_lowercase    )

       lowerCAmelCase_       :							Dict			 =		FlaxGPTaLMHeadModel.from_pretrained("""gpt2"""    )
       lowerCAmelCase_       :							List[str]			 =		False
       lowerCAmelCase_       :							List[str]			 =		model.config.eos_token_id

       lowerCAmelCase_       :							Union[str, Any]			 =		jax.jit(model.generate    )

       lowerCAmelCase_       :							int			 =		jit_generate(inputs["""input_ids"""]     ,    attention_mask=inputs["""attention_mask"""]    ).sequences

       lowerCAmelCase_       :							Optional[int]			 =		tokenizer.batch_decode(_lowercase     ,    skip_special_tokens=_lowercase    )

       lowerCAmelCase_       :							Dict			 =		[
           "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(_lowercase     ,    _lowercase    )







  @is_pt_flax_cross_test
  def 						UpperCAmelCase__		(						self    ):
       lowerCAmelCase_       :							Optional[Any]			 =		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
                 lowerCAmelCase_       :							Tuple			 =		self._prepare_for_class(_lowercase     ,    _lowercase    )
                 lowerCAmelCase_       :							Optional[int]			 =		{k: torch.tensor(v.tolist()    ) for k, v in prepared_inputs_dict.items()}

                 # load corresponding PyTorch class
                 lowerCAmelCase_       :							Optional[Any]			 =		model_class.__name__[4:]  # Skip the "Flax" at the beginning
                 lowerCAmelCase_       :							List[str]			 =		getattr(_lowercase     ,    _lowercase    )

                 lowerCAmelCase_       :							List[Any]			 =		pt_inputs["input_ids"].shape
                 lowerCAmelCase_       :							int			 =		np.random.randint(0     ,    seq_length - 1     ,    size=(batch_size,)    )
                 for batch_idx, start_index in enumerate(_lowercase    ):
                      lowerCAmelCase_       :							str			 =		0
                      lowerCAmelCase_       :							Dict			 =		1
                      lowerCAmelCase_       :							List[str]			 =		0
                      lowerCAmelCase_       :							Dict			 =		1
                 lowerCAmelCase_       :							Any			 =		pt_model_class(_lowercase    ).eval()
                 lowerCAmelCase_       :							Optional[int]			 =		model_class(_lowercase     ,    dtype=jnp.floataa    )

                 lowerCAmelCase_       :							Optional[int]			 =		convert_pytorch_state_dict_to_flax(pt_model.state_dict()     ,    _lowercase    )
                 lowerCAmelCase_       :							int			 =		fx_state

                 with torch.no_grad():
                      lowerCAmelCase_       :							Optional[int]			 =		pt_model(**_lowercase    ).to_tuple()

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

                 with tempfile.TemporaryDirectory() as tmpdirname:
                      pt_model.save_pretrained(_lowercase    )
                      lowerCAmelCase_       :							Tuple			 =		model_class.from_pretrained(_lowercase     ,    from_pt=_lowercase    )

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







  @is_pt_flax_cross_test
  def 						UpperCAmelCase__		(						self    ):
       lowerCAmelCase_       :							Tuple			 =		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
                 lowerCAmelCase_       :							Optional[Any]			 =		self._prepare_for_class(_lowercase     ,    _lowercase    )
                 lowerCAmelCase_       :							Union[str, Any]			 =		{k: torch.tensor(v.tolist()    ) for k, v in prepared_inputs_dict.items()}

                 # load corresponding PyTorch class
                 lowerCAmelCase_       :							Optional[Any]			 =		model_class.__name__[4:]  # Skip the "Flax" at the beginning
                 lowerCAmelCase_       :							int			 =		getattr(_lowercase     ,    _lowercase    )

                 lowerCAmelCase_       :							Union[str, Any]			 =		pt_model_class(_lowercase    ).eval()
                 lowerCAmelCase_       :							Union[str, Any]			 =		model_class(_lowercase     ,    dtype=jnp.floataa    )

                 lowerCAmelCase_       :							Optional[int]			 =		load_flax_weights_in_pytorch_model(_lowercase     ,    fx_model.params    )
                 lowerCAmelCase_       :							Tuple			 =		pt_inputs["input_ids"].shape
                 lowerCAmelCase_       :							Optional[Any]			 =		np.random.randint(0     ,    seq_length - 1     ,    size=(batch_size,)    )
                 for batch_idx, start_index in enumerate(_lowercase    ):
                      lowerCAmelCase_       :							str			 =		0
                      lowerCAmelCase_       :							Tuple			 =		1
                      lowerCAmelCase_       :							Union[str, Any]			 =		0
                      lowerCAmelCase_       :							List[Any]			 =		1

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

                 with torch.no_grad():
                      lowerCAmelCase_       :							str			 =		pt_model(**_lowercase    ).to_tuple()

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

                 with tempfile.TemporaryDirectory() as tmpdirname:
                      fx_model.save_pretrained(_lowercase    )
                      lowerCAmelCase_       :							Optional[Any]			 =		pt_model_class.from_pretrained(_lowercase     ,    from_flax=_lowercase    )

                 with torch.no_grad():
                      lowerCAmelCase_       :							Union[str, Any]			 =		pt_model_loaded(**_lowercase    ).to_tuple()

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







  @slow
  def 						UpperCAmelCase__		(						self    ):
       for model_class_name in self.all_model_classes:
            lowerCAmelCase_       :							Optional[int]			 =		model_class_name.from_pretrained("""gpt2"""     ,    from_pt=_lowercase    )
            lowerCAmelCase_       :							str			 =		model(np.ones((1, 1)    )    )
            self.assertIsNotNone(_lowercase    )





import inspect
import warnings

import torch

from .state import AcceleratorState, GradientState
from .utils import DistributedType, honor_type, is_tpu_available


if is_tpu_available(check_device=False):
      import torch_xla.core.xla_model as xm






def 						_lowerCAmelCase    (					__lowerCamelCase :  Any		,						__lowerCamelCase :  str  ):



 """simple docstring"""





 if isinstance(__lowerCamelCase		,						(list, tuple)  ):
  return honor_type(__lowerCamelCase		,						(move_to_device(__lowerCamelCase		,						__lowerCamelCase  ) for t in state)  )
 elif isinstance(__lowerCamelCase		,						__lowerCamelCase  ):
  return type(__lowerCamelCase  )({k: move_to_device(__lowerCamelCase		,						__lowerCamelCase  ) for k, v in state.items()}  )
 elif isinstance(__lowerCamelCase		,						torch.Tensor  ):
  return state.to(__lowerCamelCase  )
 return state






class    _SCREAMING_SNAKE_CASE						(torch.optim.Optimizer			):





     def __init__(							self			:				int  ,	UpperCamelCase			:				Any  ,	UpperCamelCase			:				List[str]=True  ,	UpperCamelCase			:				List[Any]=None					)->Optional[int]:
      __SCREAMING_SNAKE_CASE							:     str							=						optimizer
      __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						scaler
      __SCREAMING_SNAKE_CASE							:     Any							=						AcceleratorState()
      __SCREAMING_SNAKE_CASE							:     int							=						GradientState()
      __SCREAMING_SNAKE_CASE							:     str							=						device_placement
      __SCREAMING_SNAKE_CASE							:     int							=						False
      __SCREAMING_SNAKE_CASE							:     Dict							=						None

      # Handle device placement
      if device_placement:
       __SCREAMING_SNAKE_CASE							:     List[str]							=						self.optimizer.state_dict()
       if self.accelerator_state.distributed_type == DistributedType.TPU:
        xm.send_cpu_data_to_device(UpperCamelCase  ,	self.accelerator_state.device					)
       else:
        __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						move_to_device(UpperCamelCase  ,	self.accelerator_state.device					)
       self.optimizer.load_state_dict(UpperCamelCase					)





     @property
     def 						__snake_case     (							self			:				Optional[Any]					)->List[str]:
      return self.optimizer.state





     @state.setter
     def 						__snake_case     (							self			:				Optional[int]  ,	UpperCamelCase			:				List[str]					)->Dict:
      __SCREAMING_SNAKE_CASE							:     List[str]							=						state





     @property
     def 						__snake_case     (							self			:				List[Any]					)->str:
      return self.optimizer.param_groups





     @param_groups.setter
     def 						__snake_case     (							self			:				List[str]  ,	UpperCamelCase			:				Any					)->Dict:
      __SCREAMING_SNAKE_CASE							:     Optional[int]							=						param_groups





     @property
     def 						__snake_case     (							self			:				str					)->str:
      return self.optimizer.defaults





     @defaults.setter
     def 						__snake_case     (							self			:				int  ,	UpperCamelCase			:				Tuple					)->List[str]:
      __SCREAMING_SNAKE_CASE							:     List[str]							=						defaults





     def 						__snake_case     (							self			:				Any  ,	UpperCamelCase			:				Tuple					)->Optional[int]:
      self.optimizer.add_param_group(UpperCamelCase					)





     def 						__snake_case     (							self			:				Union[str, Any]  ,	UpperCamelCase			:				Union[str, Any]					)->Dict:
      if self.accelerator_state.distributed_type == DistributedType.TPU and self.device_placement:
       xm.send_cpu_data_to_device(UpperCamelCase  ,	self.accelerator_state.device					)
      self.optimizer.load_state_dict(UpperCamelCase					)





     def 						__snake_case     (							self			:				List[str]					)->List[Any]:
      return self.optimizer.state_dict()





     def 						__snake_case     (							self			:				Optional[int]  ,	UpperCamelCase			:				Any=None					)->Any:
      if self.gradient_state.sync_gradients:
       __SCREAMING_SNAKE_CASE							:     Optional[int]							=						"set_to_none" in inspect.signature(self.optimizer.zero_grad					).parameters
       if accept_arg:
        if set_to_none is None:
         __SCREAMING_SNAKE_CASE							:     Optional[int]							=						False
        self.optimizer.zero_grad(set_to_none=UpperCamelCase					)
       else:
        if set_to_none is not None:
         raise ValueError("`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer."					)
        self.optimizer.zero_grad()





     def 						__snake_case     (							self			:				Dict  ,	UpperCamelCase			:				Tuple=None					)->List[Any]:
      if self.gradient_state.sync_gradients:
       if self.accelerator_state.distributed_type == DistributedType.TPU:
        __SCREAMING_SNAKE_CASE							:     List[str]							=						{"closure": closure} if closure is not None else {}
        xm.optimizer_step(self.optimizer  ,	optimizer_args=UpperCamelCase					)
       elif self.scaler is not None:
        __SCREAMING_SNAKE_CASE							:     Optional[int]							=						False
        if self._last_scale is None:
         # `get_scale` is an async operation requiring full synchronization
         # on CPU and GPUs before finishing. As a result, we store away
         # the prior one to reduce the call overhead
         __SCREAMING_SNAKE_CASE							:     List[Any]							=						self.scaler.get_scale()
         __SCREAMING_SNAKE_CASE							:     List[Any]							=						True
        self.scaler.step(self.optimizer  ,	UpperCamelCase					)
        self.scaler.update()
        __SCREAMING_SNAKE_CASE							:     Dict							=						self.scaler.get_scale()
        if not new_scale:
         # If we reduced the loss scale, it means the optimizer step was skipped because of gradient overflow.
         __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						scale_after < self._last_scale
        __SCREAMING_SNAKE_CASE							:     List[Any]							=						scale_after

       else:
        self.optimizer.step(UpperCamelCase					)





     def 						__snake_case     (							self			:				Optional[int]  ,	UpperCamelCase			:				Optional[Any]					)->Any:
      for param_group in self.optimizer.param_groups:
       __SCREAMING_SNAKE_CASE							:     Tuple							=						[parameters_map.get(UpperCamelCase  ,	UpperCamelCase					) for p in param_group["params"]]





     @property
     def 						__snake_case     (							self			:				List[Any]					)->int:
      warnings.warn(
          "The `is_overflow` property is deprecated and will be removed in version 1.0 of Accelerate use "
          "`optimizer.step_was_skipped` instead."  ,	UpperCamelCase  ,	)
      return self._is_overflow





     @property
     def 						__snake_case     (							self			:				Optional[int]					)->List[Any]:
      return self._is_overflow





     def __getstate__(							self			:				List[str]					)->List[str]:
      return self.__dict__.copy()






     def __setstate__(							self			:				Union[str, Any]  ,	UpperCamelCase			:				str					)->int:
      self.__dict__.update(UpperCamelCase					)






