
import collections

from .utils import ExplicitEnum, is_torch_available, logging


if is_torch_available():
					import torch


lowerCamelCase__        =						logging.get_logger(__name__)







class 	A__				:






				def __init__(		self   : str       ,	a   : List[str]       ,	a   : Optional[int]=21       ,	a   : Tuple=[]       ,	a   : Optional[Any]=None  ):






							'''simple docstring'''







							lowerCAmelCase__ : List[str]   =						model
							lowerCAmelCase__ : Optional[int]   =						trace_batch_nums
							lowerCAmelCase__ : Any   =						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
							lowerCAmelCase__ : Optional[int]   =						collections.deque([]       ,	UpperCAmelCase_  )
							lowerCAmelCase__ : Optional[int]   =						[]
							lowerCAmelCase__ : Any   =						0
							lowerCAmelCase__ : Optional[int]   =						0
							lowerCAmelCase__ : Optional[int]   =						False
							lowerCAmelCase__ : str   =						'                 '

							self.analyse_model()

							self.register_forward_hook()






				def 				_lowerCamelCase						(		self   : Dict       ,	a   : int=None  ):






							'''simple docstring'''







							if frame is not None:
										self.expand_frame(UpperCAmelCase_  )
							self.frames.append('\n'.join(self.frame  )  )
							lowerCAmelCase__ : int   =						[]  # start a new frame






				def 				_lowerCamelCase						(		self   : Tuple       ,	a   : Dict  ):






							'''simple docstring'''







							self.frame.append(UpperCAmelCase_  )






				def 				_lowerCamelCase						(		self   : Optional[Any]  ):






							'''simple docstring'''







							print('\n'.join(self.frames  )  )
							lowerCAmelCase__ : Dict   =						[]






				def 				_lowerCamelCase						(		self   : Tuple  ):






							'''simple docstring'''







							lowerCAmelCase__ : Optional[int]   =						[]






				def 				_lowerCamelCase						(		self   : Any  ):






							'''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'  )
							lowerCAmelCase__ : int   =						[]






				def 				_lowerCamelCase						(		self   : Tuple  ):






							'''simple docstring'''







							lowerCAmelCase__ : 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 				_lowerCamelCase						(		self   : Any       ,	a   : int       ,	a   : List[str]  ):






							'''simple docstring'''







							if torch.is_tensor(UpperCAmelCase_  ):
										self.expand_frame(get_abs_min_max(UpperCAmelCase_       ,	UpperCAmelCase_  )  )
										if detect_overflow(UpperCAmelCase_       ,	UpperCAmelCase_  ):
													lowerCAmelCase__ : Tuple   =						True
							elif var is None:
										self.expand_frame(f'''{"None":>17} {ctx}'''  )
							else:
										self.expand_frame(f'''{"not a tensor":>17} {ctx}'''  )






				def 				_lowerCamelCase						(		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 				_lowerCamelCase						(		self   : List[str]  ):






							'''simple docstring'''







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






				def 				_lowerCamelCase						(		self   : Union[str, Any]       ,	a   : Optional[Any]       ,	a   : str       ,	a   : 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 				_lowerCamelCase						(		self   : Dict  ):






							'''simple docstring'''







							self.model.apply(self._register_forward_hook  )






				def 				_lowerCamelCase						(		self   : Any       ,	a   : str  ):






							'''simple docstring'''







							module.register_forward_hook(self.forward_hook  )






				def 				_lowerCamelCase						(		self   : Union[str, Any]       ,	a   : Optional[Any]       ,	a   : List[Any]       ,	a   : Tuple  ):






							'''simple docstring'''







							lowerCAmelCase__ : List[Any]   =						False

							lowerCAmelCase__ : Union[str, Any]   =						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
										lowerCAmelCase__ : Dict   =						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__     (    SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_	)							->						Optional[Any]:
			lowerCAmelCase__ : str   =						var.abs()
			return F'''{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}'''
def 					lowerCAmelCase__     (    SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_	)							->						Optional[int]:
			lowerCAmelCase__ : Any   =						False
			if torch.isnan(_UpperCAmelCase	).any().item():
						lowerCAmelCase__ : List[Any]   =						True
						print(F'''{ctx} has nans'''	)
			if torch.isinf(_UpperCAmelCase	).any().item():
						lowerCAmelCase__ : str   =						True
						print(F'''{ctx} has infs'''	)

			# if needed to monitor large elements can enable the following
			if 0:  # and detected:
						lowerCAmelCase__ : str   =						var[torch.ge(var.abs()			,				100	)]
						if naaa.numel() > 0:
									print(F'''{ctx}:  n100={naaa.numel()}'''	)
						lowerCAmelCase__ : int   =						var[torch.ge(var.abs()			,				1_000	)]
						if naaaa.numel() > 0:
									print(F'''{ctx}: n1000={naaaa.numel()}'''	)
						lowerCAmelCase__ : int   =						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 	A__				(	__magic_name__      ):
				lowercase            =    'underflow_overflow'
				lowercase            =    'tpu_metrics_debug'






import argparse
import os


UpperCAmelCase__    :      Any       					={
    '''pytorch''': {
        '''1.13''': {
            '''torch''': '''1.13.1''',
            '''torchvision''': '''0.14.1''',
            '''torchaudio''': '''0.13.1''',
            '''python''': 3.9,
            '''cuda''': '''cu116''',
            '''install''': (
                '''python3 -m pip install --no-cache-dir -U torch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1'''
                ''' --extra-index-url https://download.pytorch.org/whl/cu116'''
            ),
            '''base_image''': '''nvidia/cuda:11.6.2-cudnn8-devel-ubuntu20.04''',
        },
        '''1.12''': {
            '''torch''': '''1.12.1''',
            '''torchvision''': '''0.13.1''',
            '''torchaudio''': '''0.12.1''',
            '''python''': 3.9,
            '''cuda''': '''cu113''',
            '''install''': (
                '''python3 -m pip install --no-cache-dir -U torch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1'''
                ''' --extra-index-url https://download.pytorch.org/whl/cu113'''
            ),
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''1.11''': {
            '''torch''': '''1.11.0''',
            '''torchvision''': '''0.12.0''',
            '''torchaudio''': '''0.11.0''',
            '''python''': 3.9,
            '''cuda''': '''cu113''',
            '''install''': (
                '''python3 -m pip install --no-cache-dir -U torch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0'''
                ''' --extra-index-url https://download.pytorch.org/whl/cu113'''
            ),
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''1.10''': {
            '''torch''': '''1.10.2''',
            '''torchvision''': '''0.11.3''',
            '''torchaudio''': '''0.10.2''',
            '''python''': 3.9,
            '''cuda''': '''cu113''',
            '''install''': (
                '''python3 -m pip install --no-cache-dir -U torch==1.10.2 torchvision==0.11.3 torchaudio==0.10.2'''
                ''' --extra-index-url https://download.pytorch.org/whl/cu113'''
            ),
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        # torchaudio < 0.10 has no CUDA-enabled binary distributions
        '''1.9''': {
            '''torch''': '''1.9.1''',
            '''torchvision''': '''0.10.1''',
            '''torchaudio''': '''0.9.1''',
            '''python''': 3.9,
            '''cuda''': '''cu111''',
            '''install''': (
                '''python3 -m pip install --no-cache-dir -U torch==1.9.1 torchvision==0.10.1 torchaudio==0.9.1'''
                ''' --extra-index-url https://download.pytorch.org/whl/cu111'''
            ),
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
    },
    '''tensorflow''': {
        '''2.11''': {
            '''tensorflow''': '''2.11.1''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.11.1''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.10''': {
            '''tensorflow''': '''2.10.1''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.10.1''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.9''': {
            '''tensorflow''': '''2.9.3''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.9.3''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.8''': {
            '''tensorflow''': '''2.8.2''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.8.2''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.7''': {
            '''tensorflow''': '''2.7.3''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.7.3''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.6''': {
            '''tensorflow''': '''2.6.5''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.6.5''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.5''': {
            '''tensorflow''': '''2.5.3''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.5.3''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
    },
}


if __name__ == "__main__":
							UpperCAmelCase__    :      Tuple       					=argparse.ArgumentParser('''Choose the framework and version to install''')
							parser.add_argument(
							    '''--framework''', help='''The framework to install. Should be `torch` or `tensorflow`''', type=str, required=True
							)
							parser.add_argument('''--version''', help='''The version of the framework to install.''', type=str, required=True)
							UpperCAmelCase__    :      List[str]       					=parser.parse_args()

							UpperCAmelCase__    :      Dict       					=past_versions_testing[args.framework][args.version]

							os.system(F"echo \"export INSTALL_CMD=\'{info['install']}\'\" >> ~/.profile")
							print(F"echo \"export INSTALL_CMD=\'{info['install']}\'\" >> ~/.profile")

							UpperCAmelCase__    :      List[Any]       					=''''''
							if args.framework == "pytorch":
														UpperCAmelCase__    :      Optional[int]       					=info['''cuda''']
							os.system(F"echo \"export CUDA='{cuda}'\" >> ~/.profile")
							print(F"echo \"export CUDA='{cuda}'\" >> ~/.profile")





