import collections

from .utils import ExplicitEnum, is_torch_available, logging


if is_torch_available():
   import torch


lowerCamelCase__   = logging.get_logger(__name__)




class  _UpperCAmelCase       :



 '''simple docstring'''




 def __init__(						self							:					Dict				,      lowercase_							:					Optional[int]				,      lowercase_							:					Dict=21				,      lowercase_							:					Dict=[]				,      lowercase_							:					List[str]=None)		->     Union[str, Any]:


   """simple docstring"""







   _UpperCamelCase				=  model
   _UpperCamelCase				=  trace_batch_nums
   _UpperCamelCase				=  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
   _UpperCamelCase				=  collections.deque([]				,      lowercase_)
   _UpperCamelCase				=  []
   _UpperCamelCase				=  0
   _UpperCamelCase				=  0
   _UpperCamelCase				=  False
   _UpperCamelCase				=  "                 "

   self.analyse_model()

   self.register_forward_hook()







 def 						__UpperCAmelCase							(						self							:					Union[str, Any]				,      lowercase_							:					Union[str, Any]=None)		->     List[Any]:


   """simple docstring"""







   if frame is not None:
     self.expand_frame(lowercase_)
   self.frames.append("\n".join(self.frame))
   _UpperCamelCase				=  []  # start a new frame







 def 						__UpperCAmelCase							(						self							:					str				,      lowercase_							:					Tuple)		->     Dict:


   """simple docstring"""







   self.frame.append(lowercase_)







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


   """simple docstring"""







   print("\n".join(self.frames))
   _UpperCamelCase				=  []







 def 						__UpperCAmelCase							(						self							:					Optional[Any])		->     str:


   """simple docstring"""







   _UpperCamelCase				=  []







 def 						__UpperCAmelCase							(						self							:					List[Any])		->     Dict:


   """simple docstring"""







   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")
   _UpperCamelCase				=  []







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


   """simple docstring"""







   _UpperCamelCase				=  {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							:					Dict				,      lowercase_							:					Dict				,      lowercase_							:					Tuple)		->     int:


   """simple docstring"""







   if torch.is_tensor(lowercase_):
     self.expand_frame(get_abs_min_max(lowercase_				,      lowercase_))
     if detect_overflow(lowercase_				,      lowercase_):
       _UpperCamelCase				=  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							:					Optional[int])		->     Union[str, Any]:


   """simple docstring"""







   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							:					Union[str, Any])		->     Union[str, Any]:


   """simple docstring"""







   self.expand_frame(f'{self.prefix} *** Finished batch number={self.batch_number-1} ***\n\n')







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


   """simple docstring"""







   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							:					Optional[int])		->     Any:


   """simple docstring"""







   self.model.apply(self._register_forward_hook)







 def 						__UpperCAmelCase							(						self							:					str				,      lowercase_							:					List[str])		->     str:


   """simple docstring"""







   module.register_forward_hook(self.forward_hook)







 def 						__UpperCAmelCase							(						self							:					Union[str, Any]				,      lowercase_							:					int				,      lowercase_							:					int				,      lowercase_							:					Any)		->     Dict:


   """simple docstring"""







   _UpperCamelCase				=  False

   _UpperCamelCase				=  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
     _UpperCamelCase				=  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 							lowerCAmelCase__    (       a__  ,			a__      )	->Tuple:




  '''simple docstring'''

  _UpperCamelCase				=  var.abs()
  return f'{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}'


def 							lowerCAmelCase__    (       a__  ,			a__      )	->Optional[Any]:




  '''simple docstring'''

  _UpperCamelCase				=  False
  if torch.isnan(a__      ).any().item():
    _UpperCamelCase				=  True
    print(f'{ctx} has nans'      )
  if torch.isinf(a__      ).any().item():
    _UpperCamelCase				=  True
    print(f'{ctx} has infs'      )

  # if needed to monitor large elements can enable the following
  if 0:  # and detected:
    _UpperCamelCase				=  var[torch.ge(var.abs()  ,			100      )]
    if naaa.numel() > 0:
      print(f'{ctx}:  n100={naaa.numel()}'      )
    _UpperCamelCase				=  var[torch.ge(var.abs()  ,			1_000      )]
    if naaaa.numel() > 0:
      print(f'{ctx}: n1000={naaaa.numel()}'      )
    _UpperCamelCase				=  var[torch.ge(var.abs()  ,			10_000      )]
    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       (			lowerCAmelCase   ):



 '''simple docstring'''




 __A           = '''underflow_overflow'''
 __A           = '''tpu_metrics_debug'''





lowerCamelCase__   = '''
# Transformers installation
! pip install transformers datasets
# To install from source instead of the last release, comment the command above and uncomment the following one.
# ! pip install git+https://github.com/huggingface/transformers.git
'''

lowerCamelCase__   = [{'''type''': '''code''', '''content''': INSTALL_CONTENT}]
lowerCamelCase__   = {
    '''{processor_class}''': '''FakeProcessorClass''',
    '''{model_class}''': '''FakeModelClass''',
    '''{object_class}''': '''FakeObjectClass''',
}




