
"""simple docstring"""




import collections

from .utils import ExplicitEnum, is_torch_available, logging


if is_torch_available():
					import torch


UpperCAmelCase__				:   Optional[Any]   							=			logging.get_logger(__name__)



class      lowerCAmelCase_						:


				"""simple docstring"""



				def __init__(self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__=21       ,	SCREAMING_SNAKE_CASE__=[]       ,	SCREAMING_SNAKE_CASE__=None   ) ->		Optional[Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Dict									=       model
											SCREAMING_SNAKE_CASE__						:	Tuple									=       trace_batch_nums
											SCREAMING_SNAKE_CASE__						:	Tuple									=       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
											SCREAMING_SNAKE_CASE__						:	List[str]									=       collections.deque([]       ,	SCREAMING_SNAKE_CASE__   )
											SCREAMING_SNAKE_CASE__						:	Tuple									=       []
											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       0
											SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       0
											SCREAMING_SNAKE_CASE__						:	Optional[int]									=       False
											SCREAMING_SNAKE_CASE__						:	List[Any]									=       """                 """

											self.analyse_model()

											self.register_forward_hook()



				def      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__=None   ) ->		Tuple:



											"""simple docstring"""




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



				def      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__   ) ->		int:



											"""simple docstring"""




											self.frame.append(SCREAMING_SNAKE_CASE__   )



				def      __magic_name__   (self   ) ->		Any:



											"""simple docstring"""




											print("""\n""".join(self.frames   )   )
											SCREAMING_SNAKE_CASE__						:	Optional[Any]									=       []



				def      __magic_name__   (self   ) ->		List[Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	int									=       []



				def      __magic_name__   (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"""   )
											SCREAMING_SNAKE_CASE__						:	List[Any]									=       []



				def      __magic_name__   (self   ) ->		Any:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	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      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ) ->		str:



											"""simple docstring"""




											if torch.is_tensor(SCREAMING_SNAKE_CASE__   ):
																		self.expand_frame(get_abs_min_max(SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   )   )
																		if detect_overflow(SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ):
																									SCREAMING_SNAKE_CASE__						:	str									=       True
											elif var is None:
																		self.expand_frame(F'''{'None':>17} {ctx}'''   )
											else:
																		self.expand_frame(F'''{'not a tensor':>17} {ctx}'''   )



				def      __magic_name__   (self   ) ->		Tuple:



											"""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      __magic_name__   (self   ) ->		int:



											"""simple docstring"""




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



				def      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ) ->		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=SCREAMING_SNAKE_CASE__   ):
																		self.analyse_variable(SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   )

											# inputs
											if isinstance(SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ):
																		for i, x in enumerate(SCREAMING_SNAKE_CASE__   ):
																									self.analyse_variable(SCREAMING_SNAKE_CASE__       ,	F'''input[{i}]'''   )
											else:
																		self.analyse_variable(SCREAMING_SNAKE_CASE__       ,	"""input"""   )

											# outputs
											if isinstance(SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ):
																		for i, x in enumerate(SCREAMING_SNAKE_CASE__   ):
																									# possibly a tuple of tuples
																									if isinstance(SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ):
																																for j, y in enumerate(SCREAMING_SNAKE_CASE__   ):
																																							self.analyse_variable(SCREAMING_SNAKE_CASE__       ,	F'''output[{i}][{j}]'''   )
																									else:
																																self.analyse_variable(SCREAMING_SNAKE_CASE__       ,	F'''output[{i}]'''   )
											else:
																		self.analyse_variable(SCREAMING_SNAKE_CASE__       ,	"""output"""   )

											self.save_frame()



				def      __magic_name__   (self   ) ->		str:



											"""simple docstring"""




											self.model.apply(self._register_forward_hook   )



				def      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__   ) ->		str:



											"""simple docstring"""




											module.register_forward_hook(self.forward_hook   )



				def      __magic_name__   (self       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   ) ->		int:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__						:	Dict									=       False

											SCREAMING_SNAKE_CASE__						:	str									=       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
																		SCREAMING_SNAKE_CASE__						:	List[Any]									=       True

											self.create_frame(SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__       ,	SCREAMING_SNAKE_CASE__   )

											# 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 lowercase_  (       _snake_case       ,_snake_case						):
							SCREAMING_SNAKE_CASE__						:	Any									=       var.abs()
							return f'''{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}'''

def lowercase_  (       _snake_case       ,_snake_case						):
							SCREAMING_SNAKE_CASE__						:	Union[str, Any]									=       False
							if torch.isnan(_snake_case						).any().item():
														SCREAMING_SNAKE_CASE__						:	str									=       True
														print(f'''{ctx} has nans'''						)
							if torch.isinf(_snake_case						).any().item():
														SCREAMING_SNAKE_CASE__						:	Optional[int]									=       True
														print(f'''{ctx} has infs'''						)

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


				"""simple docstring"""



				__UpperCamelCase				:						str 	=       '''underflow_overflow'''
				__UpperCamelCase				:						Union[str, Any] 	=       '''tpu_metrics_debug'''






"""simple docstring"""




import re

def lowercase_  (       _snake_case						):
							def docstring_decorator(_snake_case						):
														SCREAMING_SNAKE_CASE__						:	Any									=       fn.__doc__
														SCREAMING_SNAKE_CASE__						:	Dict									=       func_doc.split("""\n"""						)
														SCREAMING_SNAKE_CASE__						:	Any									=       0
														while i < len(_snake_case						) and re.search(R"""^\s*Examples?:\s*$"""       ,lines[i]						) is None:
																					i += 1
														if i < len(_snake_case						):
																					SCREAMING_SNAKE_CASE__						:	List[str]									=       example_docstring
																					SCREAMING_SNAKE_CASE__						:	List[Any]									=       """\n""".join(_snake_case						)
														else:
																					raise ValueError(
																					    f'''The function {fn} should have an empty \'Examples:\' in its docstring as placeholder, '''
																					    f'''current docstring is:\n{func_doc}'''						)
														SCREAMING_SNAKE_CASE__						:	int									=       func_doc
														return fn

							return docstring_decorator




