code: 
  
  
  
  
  
  import collections
  
  from .utils import ExplicitEnum, is_torch_available, logging
  
  
  if is_torch_available():
       import torch
  
  
  _UpperCamelCase       = logging.get_logger(__name__)
  
  
  
  
  class 				_lowerCamelCase :
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        def __init__( self				,						UpperCAmelCase				,						UpperCAmelCase=21				,						UpperCAmelCase=[]				,						UpperCAmelCase=None   )       ->   Any:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				Dict        =       model
          __snake_case   :				Dict        =       trace_batch_nums
          __snake_case   :				str        =       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
          __snake_case   :				Any        =       collections.deque([]				,						UpperCAmelCase   )
          __snake_case   :				Optional[int]        =       []
          __snake_case   :				List[Any]        =       0
          __snake_case   :				List[Any]        =       0
          __snake_case   :				str        =       False
          __snake_case   :				str        =       "                 "
  
          self.analyse_model()
  
          self.register_forward_hook()
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self				,						UpperCAmelCase=None   )       ->   List[Any]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          if frame is not None:
            self.expand_frame(UpperCAmelCase   )
          self.frames.append("\n".join(self.frame   )   )
          __snake_case   :				int        =       []  # start a new frame
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self				,						UpperCAmelCase   )       ->   List[Any]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          self.frame.append(UpperCAmelCase   )
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   Any:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          print("\n".join(self.frames   )   )
          __snake_case   :				Union[str, Any]        =       []
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   Optional[Any]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				Dict        =       []
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   int:
  
  
  
  
  
  
          '''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"   )
          __snake_case   :				Union[str, Any]        =       []
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   Any:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				List[Any]        =       {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				,						UpperCAmelCase				,						UpperCAmelCase   )       ->   Tuple:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          if torch.is_tensor(UpperCAmelCase   ):
            self.expand_frame(get_abs_min_max(UpperCAmelCase				,						UpperCAmelCase   )   )
            if detect_overflow(UpperCAmelCase				,						UpperCAmelCase   ):
              __snake_case   :				List[Any]        =       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   )       ->   int:
  
  
  
  
  
  
          '''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   )       ->   List[Any]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          self.expand_frame(F"""{self.prefix} *** Finished batch number={self.batch_number-1} ***\n\n"""   )
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase   )       ->   str:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          self.expand_frame(F"""{self.prefix} {self.module_names[module]} {module.__class__.__name__}"""   )
  
          # params
          for name, p in module.named_parameters(recurse=UpperCAmelCase   ):
            self.analyse_variable(UpperCAmelCase				,						UpperCAmelCase   )
  
          # inputs
          if isinstance(UpperCAmelCase				,						UpperCAmelCase   ):
            for i, x in enumerate(UpperCAmelCase   ):
              self.analyse_variable(UpperCAmelCase				,						F"""input[{i}]"""   )
          else:
            self.analyse_variable(UpperCAmelCase				,						"input"   )
  
          # outputs
          if isinstance(UpperCAmelCase				,						UpperCAmelCase   ):
            for i, x in enumerate(UpperCAmelCase   ):
              # possibly a tuple of tuples
              if isinstance(UpperCAmelCase				,						UpperCAmelCase   ):
                for j, y in enumerate(UpperCAmelCase   ):
                  self.analyse_variable(UpperCAmelCase				,						F"""output[{i}][{j}]"""   )
              else:
                self.analyse_variable(UpperCAmelCase				,						F"""output[{i}]"""   )
          else:
            self.analyse_variable(UpperCAmelCase				,						"output"   )
  
          self.save_frame()
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self   )       ->   Union[str, Any]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          self.model.apply(self._register_forward_hook   )
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self				,						UpperCAmelCase   )       ->   Any:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          module.register_forward_hook(self.forward_hook   )
  
  
  
  
  
  
  
        def 						UpperCAmelCase		( self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase   )       ->   List[Any]:
  
  
  
  
  
  
          '''simple docstring'''
  
  
  
  
  
  
          __snake_case   :				int        =       False
  
          __snake_case   :				int        =       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
            __snake_case   :				List[str]        =       True
  
          self.create_frame(UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase   )
  
          # 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__(	lowercase       :    int   ,     lowercase       :    Dict				) ->    Optional[int]:
    __snake_case   :				Union[str, Any]        =       var.abs()
    return f"""{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}"""
  
  
  
  
  
  
  
  def       lowerCAmelCase__(	lowercase       :    Union[str, Any]   ,     lowercase       :    str				) ->    Optional[Any]:
    __snake_case   :				List[Any]        =       False
    if torch.isnan(lowercase				).any().item():
      __snake_case   :				Any        =       True
      print(f"""{ctx} has nans"""				)
    if torch.isinf(lowercase				).any().item():
      __snake_case   :				Union[str, Any]        =       True
      print(f"""{ctx} has infs"""				)
  
    # if needed to monitor large elements can enable the following
    if 0:  # and detected:
      __snake_case   :				int        =       var[torch.ge(var.abs()   ,     100				)]
      if naaa.numel() > 0:
        print(f"""{ctx}:  n100={naaa.numel()}"""				)
      __snake_case   :				List[Any]        =       var[torch.ge(var.abs()   ,     1000				)]
      if naaaa.numel() > 0:
        print(f"""{ctx}: n1000={naaaa.numel()}"""				)
      __snake_case   :				Tuple        =       var[torch.ge(var.abs()   ,     1_0000				)]
      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 				_lowerCamelCase (    a    ):
  
  
  
        """simple docstring"""
  
  
  
  
  
  
        UpperCAmelCase_						:      List[str]       ="underflow_overflow"
        UpperCAmelCase_						:      Optional[Any]       ="tpu_metrics_debug"
  
  
style_context: 
  
  
  
  
  
  from typing import TYPE_CHECKING
  
  from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available
  
  
  _UpperCamelCase       = {'''configuration_nat''': ['''NAT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''NatConfig''']}
  
  
  try:
       if not is_torch_available():
            raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
       pass
  else:
       _UpperCamelCase       = [
           '''NAT_PRETRAINED_MODEL_ARCHIVE_LIST''',
           '''NatForImageClassification''',
           '''NatModel''',
           '''NatPreTrainedModel''',
           '''NatBackbone''',
       ]
  
  if TYPE_CHECKING:
       from .configuration_nat import NAT_PRETRAINED_CONFIG_ARCHIVE_MAP, NatConfig
  
       try:
            if not is_torch_available():
                 raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
            pass
       else:
            from .modeling_nat import (
                NAT_PRETRAINED_MODEL_ARCHIVE_LIST,
                NatBackbone,
                NatForImageClassification,
                NatModel,
                NatPreTrainedModel,
            )
  
  else:
       import sys
  
       _UpperCamelCase       = _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)
  
  