code: 
  
  
  
  
  
  '''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      UpperCAmelCase			:
   def __init__(       self   :Optional[Any]		, lowercase_   :int		, lowercase_   :Tuple=14		, lowercase_   :List[str]=7		, lowercase_   :Any=True		, lowercase_   :Dict=True		, lowercase_   :str=False		, lowercase_   :Union[str, Any]=True		, lowercase_   :List[Any]=99		, lowercase_   :int=32		, lowercase_   :Tuple=5		, lowercase_   :str=4		, lowercase_   :Union[str, Any]=37		, lowercase_   :List[Any]="gelu"		, lowercase_   :Tuple=0.1		, lowercase_   :Dict=0.1		, lowercase_   :Dict=5_12		, lowercase_   :str=0.0_2		, )->							int:
       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   UpperCAmelCase_							(       self   :Optional[Any]							)->							Dict:
       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=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   :Optional[Any]							)->							Dict:
       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   UpperCAmelCase_							(       self   :str							)->							Optional[int]:
       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   UpperCAmelCase_							(       self   :Optional[int]		, lowercase_   :Any		, lowercase_   :Any		, lowercase_   :Union[str, Any]		, lowercase_   :str							)->							Optional[int]:
       A__       =    20
       A__       =    model_class_name(lowercase_							)
  
       A__       =    model.init_cache(input_ids.shape[0]		, lowercase_							)
       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=lowercase_		, past_key_values=lowercase_		, position_ids=lowercase_		, )
  
       A__       =    jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]		, dtype="i4"							)
       A__       =    model(
           input_ids[:, -1:]		, attention_mask=lowercase_		, past_key_values=outputs_cache.past_key_values		, position_ids=lowercase_		, )
  
       A__       =    model(lowercase_							)
  
       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   UpperCAmelCase_							(       self   :List[str]		, lowercase_   :str		, lowercase_   :int		, lowercase_   :List[Any]		, lowercase_   :Tuple							)->							List[Any]:
       A__       =    20
       A__       =    model_class_name(lowercase_							)
  
       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]		, lowercase_							)
       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=lowercase_		, past_key_values=lowercase_		, position_ids=lowercase_		, )
       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=lowercase_		, position_ids=lowercase_		, )
  
       A__       =    model(lowercase_		, attention_mask=lowercase_							)
  
       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      UpperCAmelCase			( UpperCamelCase__    , UpperCamelCase__    , unittest.TestCase       ):
   __lowercase     	=       (FlaxGPTaModel, FlaxGPTaLMHeadModel) if is_flax_available() else ()
   __lowercase     	=       (FlaxGPTaLMHeadModel,) if is_flax_available() else ()
   def   UpperCAmelCase_							(       self   :Union[str, Any]							)->							Optional[Any]:
       A__       =    FlaxGPTaModelTester(self							)
   def   UpperCAmelCase_							(       self   :Optional[Any]							)->							Optional[Any]:
       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(lowercase_		, lowercase_		, lowercase_		, lowercase_							)
   def   UpperCAmelCase_							(       self   :List[Any]							)->							str:
       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(
               lowercase_		, lowercase_		, lowercase_		, lowercase_							)
   @slow
   def   UpperCAmelCase_							(       self   :Optional[Any]							)->							Optional[Any]:
       A__       =    GPTaTokenizer.from_pretrained("gpt2"		, pad_token="</s>"		, padding_side="left"							)
       A__       =    tokenizer(["Hello this is a long string", "Hey"]		, return_tensors="np"		, padding=lowercase_		, truncation=lowercase_							)
  
       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(lowercase_		, skip_special_tokens=lowercase_							)
  
       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(lowercase_		, lowercase_							)
   @is_pt_flax_cross_test
   def   UpperCAmelCase_							(       self   :Any							)->							int:
       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(lowercase_		, lowercase_							)
               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(lowercase_		, lowercase_							)
  
               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(lowercase_							):
                   A__       =    0
                   A__       =    1
                   A__       =    0
                   A__       =    1
               A__       =    pt_model_class(lowercase_							).eval()
               A__       =    model_class(lowercase_		, dtype=jnp.floataa							)
  
               A__       =    convert_pytorch_state_dict_to_flax(pt_model.state_dict()		, lowercase_							)
               A__       =    fx_state
  
               with torch.no_grad():
                   A__       =    pt_model(**lowercase_							).to_tuple()
  
               A__       =    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_							)
                   A__       =    model_class.from_pretrained(lowercase_		, from_pt=lowercase_							)
  
               A__       =    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   :List[str]							)->							int:
       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(lowercase_		, lowercase_							)
               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(lowercase_		, lowercase_							)
  
               A__       =    pt_model_class(lowercase_							).eval()
               A__       =    model_class(lowercase_		, dtype=jnp.floataa							)
  
               A__       =    load_flax_weights_in_pytorch_model(lowercase_		, 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(lowercase_							):
                   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(**lowercase_							).to_tuple()
  
               A__       =    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_							)
                   A__       =    pt_model_class.from_pretrained(lowercase_		, from_flax=lowercase_							)
  
               with torch.no_grad():
                   A__       =    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   :List[str]							)->							str:
       for model_class_name in self.all_model_classes:
           A__       =    model_class_name.from_pretrained("gpt2"		, from_pt=lowercase_							)
           A__       =    model(np.ones((1, 1)							)							)
           self.assertIsNotNone(lowercase_							)
  
  
  
  
style_context: 
  
  
  
  
  
  '''simple docstring'''
  
  import collections
  
  from .utils import ExplicitEnum, is_torch_available, logging
  
  
  if is_torch_available():
    import torch
  
  
  __lowerCAmelCase							:			Optional[int]				=logging.get_logger(__name__)
  
  
  
  
  
  
  class      UpperCAmelCase			:
   def __init__(       self   :Optional[int]		, lowercase_   :Any		, lowercase_   :Any=21		, lowercase_   :Any=[]		, lowercase_   :Dict=None							)->							Tuple:
       A__       =    model
       A__       =    trace_batch_nums
       A__       =    abort_after_batch_num
  
       # keep a LIFO buffer of frames to dump as soon as inf/nan is encountered to give context to the problem emergence
       A__       =    collections.deque([]		, lowercase_							)
       A__       =    []
       A__       =    0
       A__       =    0
       A__       =    False
       A__       =    "                 "
  
       self.analyse_model()
  
       self.register_forward_hook()
   def   UpperCAmelCase_							(       self   :Any		, lowercase_   :Any=None							)->							Optional[Any]:
       if frame is not None:
           self.expand_frame(lowercase_							)
       self.frames.append("\n".join(self.frame							)							)
       A__       =    []  # start a new frame
   def   UpperCAmelCase_							(       self   :Union[str, Any]		, lowercase_   :Dict							)->							List[str]:
       self.frame.append(lowercase_							)
   def   UpperCAmelCase_							(       self   :int							)->							Dict:
       print("\n".join(self.frames							)							)
       A__       =    []
   def   UpperCAmelCase_							(       self   :str							)->							Union[str, Any]:
       A__       =    []
   def   UpperCAmelCase_							(       self   :List[str]							)->							Optional[Any]:
       print(F"\nDetected inf/nan during batch_number={self.batch_number}"							)
       print(F"Last {len(self.frames							)} forward frames:"							)
       print(F"{'abs min':8} {'abs max':8} metadata"							)
       print("\n".join(self.frames							)							)
       print("\n\n"							)
       A__       =    []
   def   UpperCAmelCase_							(       self   :Any							)->							Optional[Any]:
       # extract the fully qualified module names, to be able to report at run time. e.g.:
       # encoder.block.2.layer.0.SelfAttention.o
       #
       # for shared weights only the first shared module name will be registered
       A__       =    {m: name for name, m in self.model.named_modules()}
       # self.longest_module_name = max(len(v) for v in self.module_names.values())
   def   UpperCAmelCase_							(       self   :Tuple		, lowercase_   :Dict		, lowercase_   :Any							)->							str:
       if torch.is_tensor(lowercase_							):
           self.expand_frame(get_abs_min_max(lowercase_		, lowercase_							)							)
           if detect_overflow(lowercase_		, lowercase_							):
               A__       =    True
       elif var is None:
           self.expand_frame(F"{'None':>17} {ctx}"							)
       else:
           self.expand_frame(F"{'not a tensor':>17} {ctx}"							)
   def   UpperCAmelCase_							(       self   :Union[str, Any]							)->							str:
       self.expand_frame(F"\n\n{self.prefix} *** Starting batch number={self.batch_number} ***"							)
       self.expand_frame(F"{'abs min':8} {'abs max':8} metadata"							)
   def   UpperCAmelCase_							(       self   :int							)->							int:
       self.expand_frame(F"{self.prefix} *** Finished batch number={self.batch_number-1} ***\n\n"							)
   def   UpperCAmelCase_							(       self   :Optional[int]		, lowercase_   :Dict		, lowercase_   :Dict		, lowercase_   :str							)->							int:
       self.expand_frame(F"{self.prefix} {self.module_names[module]} {module.__class__.__name__}"							)
  
       # params
       for name, p in module.named_parameters(recurse=lowercase_							):
           self.analyse_variable(lowercase_		, lowercase_							)
  
       # inputs
       if isinstance(lowercase_		, lowercase_							):
           for i, x in enumerate(lowercase_							):
               self.analyse_variable(lowercase_		, F"input[{i}]"							)
       else:
           self.analyse_variable(lowercase_		, "input"							)
  
       # outputs
       if isinstance(lowercase_		, lowercase_							):
           for i, x in enumerate(lowercase_							):
               # possibly a tuple of tuples
               if isinstance(lowercase_		, lowercase_							):
                   for j, y in enumerate(lowercase_							):
                       self.analyse_variable(lowercase_		, F"output[{i}][{j}]"							)
               else:
                   self.analyse_variable(lowercase_		, F"output[{i}]"							)
       else:
           self.analyse_variable(lowercase_		, "output"							)
  
       self.save_frame()
   def   UpperCAmelCase_							(       self   :Any							)->							int:
       self.model.apply(self._register_forward_hook							)
   def   UpperCAmelCase_							(       self   :Tuple		, lowercase_   :Optional[int]							)->							Tuple:
       module.register_forward_hook(self.forward_hook							)
  
  
  
   def   UpperCAmelCase_							(       self   :List[str]		, lowercase_   :int		, lowercase_   :Optional[Any]		, lowercase_   :Dict							)->							Union[str, Any]:
       # - input is a tuple of packed inputs (could be non-Tensors)
       # - output could be a Tensor or a tuple of Tensors and non-Tensors
  
       A__       =    False
  
       A__       =    True if self.batch_number in self.trace_batch_nums else False
       if trace_mode:
           self.reset_saved_frames()
  
       if self.total_calls == 0:
           self.batch_start_frame()
       self.total_calls += 1
  
       # count batch numbers - the very first forward hook of the batch will be called when the
       # batch completes - i.e. it gets called very last - we know this batch has finished
       if module == self.model:
           self.batch_number += 1
           A__       =    True
  
       self.create_frame(lowercase_		, lowercase_		, lowercase_							)
  
       # if last_frame_of_batch:
       #     self.batch_end_frame()
  
       if trace_mode:
           self.trace_frames()
  
       if last_frame_of_batch:
           self.batch_start_frame()
  
       if self.detected_overflow and not trace_mode:
           self.dump_saved_frames()
  
           # now we can abort, as it's pointless to continue running
           raise ValueError(
               "DebugUnderflowOverflow: inf/nan detected, aborting as there is no point running further. "
               "Please scroll up above this traceback to see the activation values prior to this event."							)
  
       # abort after certain batch if requested to do so
       if self.abort_after_batch_num is not None and self.batch_number > self.abort_after_batch_num:
           raise ValueError(
               F"DebugUnderflowOverflow: aborting after {self.batch_number} batches due to"
               F" `abort_after_batch_num={self.abort_after_batch_num}` arg"							)
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      Dict		, _lowerCamelCase      :      int ):
      A__       =    var.abs()
      return F"{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}"
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      Tuple		, _lowerCamelCase      :      Optional[int] ):
      A__       =    False
      if torch.isnan(_lowerCamelCase ).any().item():
          A__       =    True
          print(F"{ctx} has nans" )
      if torch.isinf(_lowerCamelCase ).any().item():
          A__       =    True
          print(F"{ctx} has infs" )
  
      # if needed to monitor large elements can enable the following
      if 0:  # and detected:
          A__       =    var[torch.ge(var.abs()		, 1_00 )]
          if naaa.numel() > 0:
              print(F"{ctx}:  n100={naaa.numel()}" )
          A__       =    var[torch.ge(var.abs()		, 10_00 )]
          if naaaa.numel() > 0:
              print(F"{ctx}: n1000={naaaa.numel()}" )
          A__       =    var[torch.ge(var.abs()		, 1_00_00 )]
          if naaaaa.numel() > 0:
              print(F"{ctx}: n10000={naaaaa.numel()}" )
  
      if 0:
          print(F"min={var.min():9.2e} max={var.max():9.2e}" )
  
      if 0:
          print(F"min={var.min():9.2e} max={var.max():9.2e} var={var.var():9.2e} mean={var.mean():9.2e} ({ctx})" )
  
      return detected
  
  
  
  
  
  
  class      UpperCAmelCase			( UpperCamelCase__       ):
   __lowercase     	=       """underflow_overflow"""
   __lowercase     	=       """tpu_metrics_debug"""
  
  
  
  