



'''simple docstring'''







import collections
import inspect
import unittest

from transformers import DinatConfig
from transformers.testing_utils import require_natten, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch
							from torch import nn

							from transformers import DinatBackbone, DinatForImageClassification, DinatModel
							from transformers.models.dinat.modeling_dinat import DINAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
							from PIL import Image

							from transformers import AutoImageProcessor




class 				__lowerCAmelCase:






			def __init__(   self  :							Optional[Any]      ,       SCREAMING_SNAKE_CASE  :							Optional[Any]      ,       SCREAMING_SNAKE_CASE  :							int=13      ,       SCREAMING_SNAKE_CASE  :							Any=64      ,       SCREAMING_SNAKE_CASE  :							Tuple=4      ,       SCREAMING_SNAKE_CASE  :							List[str]=3      ,       SCREAMING_SNAKE_CASE  :							List[Any]=16      ,       SCREAMING_SNAKE_CASE  :							Union[str, Any]=[1, 2, 1]      ,       SCREAMING_SNAKE_CASE  :							List[Any]=[2, 4, 8]      ,       SCREAMING_SNAKE_CASE  :							Dict=3      ,       SCREAMING_SNAKE_CASE  :							Union[str, Any]=[[3], [1, 2], [1]]      ,       SCREAMING_SNAKE_CASE  :							Any=2.0      ,       SCREAMING_SNAKE_CASE  :							str=True      ,       SCREAMING_SNAKE_CASE  :							Optional[Any]=0.0      ,       SCREAMING_SNAKE_CASE  :							Dict=0.0      ,       SCREAMING_SNAKE_CASE  :							Optional[int]=0.1      ,       SCREAMING_SNAKE_CASE  :							str="gelu"      ,       SCREAMING_SNAKE_CASE  :							List[str]=True      ,       SCREAMING_SNAKE_CASE  :							str=0.02      ,       SCREAMING_SNAKE_CASE  :							str=1E-5      ,       SCREAMING_SNAKE_CASE  :							Tuple=True      ,       SCREAMING_SNAKE_CASE  :							Tuple=None      ,       SCREAMING_SNAKE_CASE  :							Any=True      ,       SCREAMING_SNAKE_CASE  :							Any=10      ,       SCREAMING_SNAKE_CASE  :							Optional[int]=["stage1", "stage2"]      ,       SCREAMING_SNAKE_CASE  :							Optional[Any]=[1, 2]      ,       ):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      parent
					SCREAMING_SNAKE_CASE_						:List[Any]										=      batch_size
					SCREAMING_SNAKE_CASE_						:Tuple										=      image_size
					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      patch_size
					SCREAMING_SNAKE_CASE_						:Optional[int]										=      num_channels
					SCREAMING_SNAKE_CASE_						:List[str]										=      embed_dim
					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      depths
					SCREAMING_SNAKE_CASE_						:List[str]										=      num_heads
					SCREAMING_SNAKE_CASE_						:Optional[int]										=      kernel_size
					SCREAMING_SNAKE_CASE_						:Tuple										=      dilations
					SCREAMING_SNAKE_CASE_						:str										=      mlp_ratio
					SCREAMING_SNAKE_CASE_						:Dict										=      qkv_bias
					SCREAMING_SNAKE_CASE_						:str										=      hidden_dropout_prob
					SCREAMING_SNAKE_CASE_						:Any										=      attention_probs_dropout_prob
					SCREAMING_SNAKE_CASE_						:List[str]										=      drop_path_rate
					SCREAMING_SNAKE_CASE_						:Any										=      hidden_act
					SCREAMING_SNAKE_CASE_						:Any										=      patch_norm
					SCREAMING_SNAKE_CASE_						:Tuple										=      layer_norm_eps
					SCREAMING_SNAKE_CASE_						:int										=      initializer_range
					SCREAMING_SNAKE_CASE_						:Optional[int]										=      is_training
					SCREAMING_SNAKE_CASE_						:List[Any]										=      scope
					SCREAMING_SNAKE_CASE_						:List[Any]										=      use_labels
					SCREAMING_SNAKE_CASE_						:List[Any]										=      num_labels
					SCREAMING_SNAKE_CASE_						:Optional[int]										=      out_features
					SCREAMING_SNAKE_CASE_						:Dict										=      out_indices






			def 				_lowercase	(   self  :							List[Any]	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Optional[int]										=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      None
					if self.use_labels:
							SCREAMING_SNAKE_CASE_						:Dict										=      ids_tensor([self.batch_size]      ,       self.num_labels	)

					SCREAMING_SNAKE_CASE_						:List[str]										=      self.get_config()

					return config, pixel_values, labels






			def 				_lowercase	(   self  :							List[Any]	):




					"""simple docstring"""







					return DinatConfig(
					    num_labels=self.num_labels      ,       image_size=self.image_size      ,       patch_size=self.patch_size      ,       num_channels=self.num_channels      ,       embed_dim=self.embed_dim      ,       depths=self.depths      ,       num_heads=self.num_heads      ,       kernel_size=self.kernel_size      ,       dilations=self.dilations      ,       mlp_ratio=self.mlp_ratio      ,       qkv_bias=self.qkv_bias      ,       hidden_dropout_prob=self.hidden_dropout_prob      ,       attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,       drop_path_rate=self.drop_path_rate      ,       hidden_act=self.hidden_act      ,       patch_norm=self.patch_norm      ,       layer_norm_eps=self.layer_norm_eps      ,       initializer_range=self.initializer_range      ,       out_features=self.out_features      ,       out_indices=self.out_indices      ,       )






			def 				_lowercase	(   self  :							Dict      ,       SCREAMING_SNAKE_CASE  :							Union[str, Any]      ,       SCREAMING_SNAKE_CASE  :							List[str]      ,       SCREAMING_SNAKE_CASE  :							List[Any]	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:int										=      DinatModel(config=SCREAMING_SNAKE_CASE	)
					model.to(SCREAMING_SNAKE_CASE	)
					model.eval()
					SCREAMING_SNAKE_CASE_						:Tuple										=      model(SCREAMING_SNAKE_CASE	)

					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      (config.image_size // config.patch_size) // (2 ** (len(config.depths	) - 1))
					SCREAMING_SNAKE_CASE_						:Optional[int]										=      int(config.embed_dim * 2 ** (len(config.depths	) - 1)	)

					self.parent.assertEqual(
					    result.last_hidden_state.shape      ,       (self.batch_size, expected_height, expected_width, expected_dim)	)






			def 				_lowercase	(   self  :							List[Any]      ,       SCREAMING_SNAKE_CASE  :							Union[str, Any]      ,       SCREAMING_SNAKE_CASE  :							List[str]      ,       SCREAMING_SNAKE_CASE  :							Optional[Any]	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Tuple										=      DinatForImageClassification(SCREAMING_SNAKE_CASE	)
					model.to(SCREAMING_SNAKE_CASE	)
					model.eval()
					SCREAMING_SNAKE_CASE_						:str										=      model(SCREAMING_SNAKE_CASE      ,       labels=SCREAMING_SNAKE_CASE	)
					self.parent.assertEqual(result.logits.shape      ,       (self.batch_size, self.num_labels)	)

					# test greyscale images
					SCREAMING_SNAKE_CASE_						:str										=      1
					SCREAMING_SNAKE_CASE_						:Any										=      DinatForImageClassification(SCREAMING_SNAKE_CASE	)
					model.to(SCREAMING_SNAKE_CASE	)
					model.eval()

					SCREAMING_SNAKE_CASE_						:str										=      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
					SCREAMING_SNAKE_CASE_						:Tuple										=      model(SCREAMING_SNAKE_CASE	)
					self.parent.assertEqual(result.logits.shape      ,       (self.batch_size, self.num_labels)	)






			def 				_lowercase	(   self  :							str      ,       SCREAMING_SNAKE_CASE  :							Dict      ,       SCREAMING_SNAKE_CASE  :							Tuple      ,       SCREAMING_SNAKE_CASE  :							Optional[int]	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Dict										=      DinatBackbone(config=SCREAMING_SNAKE_CASE	)
					model.to(SCREAMING_SNAKE_CASE	)
					model.eval()
					SCREAMING_SNAKE_CASE_						:Optional[int]										=      model(SCREAMING_SNAKE_CASE	)

					# verify hidden states
					self.parent.assertEqual(len(result.feature_maps	)      ,       len(config.out_features	)	)
					self.parent.assertListEqual(list(result.feature_maps[0].shape	)      ,       [self.batch_size, model.channels[0], 16, 16]	)

					# verify channels
					self.parent.assertEqual(len(model.channels	)      ,       len(config.out_features	)	)

					# verify backbone works with out_features=None
					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      None
					SCREAMING_SNAKE_CASE_						:Dict										=      DinatBackbone(config=SCREAMING_SNAKE_CASE	)
					model.to(SCREAMING_SNAKE_CASE	)
					model.eval()
					SCREAMING_SNAKE_CASE_						:List[Any]										=      model(SCREAMING_SNAKE_CASE	)

					# verify feature maps
					self.parent.assertEqual(len(result.feature_maps	)      ,       1	)
					self.parent.assertListEqual(list(result.feature_maps[0].shape	)      ,       [self.batch_size, model.channels[-1], 4, 4]	)

					# verify channels
					self.parent.assertEqual(len(model.channels	)      ,       1	)






			def 				_lowercase	(   self  :							str	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      self.prepare_config_and_inputs()
					SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						:List[str]										=      config_and_inputs
					SCREAMING_SNAKE_CASE_						:Tuple										=      {'pixel_values': pixel_values}
					return config, inputs_dict




@require_natten
@require_torch
class 				__lowerCAmelCase(							lowerCAmelCase__		,			lowerCAmelCase__		,			unittest.TestCase      ):
			__snake_case  :		Dict		=      (
			    (
			        DinatModel,
			        DinatForImageClassification,
			        DinatBackbone,
			    )
			    if is_torch_available()
			    else ()
			)
			__snake_case  :		str		=      (
			    {'feature-extraction': DinatModel, 'image-classification': DinatForImageClassification}
			    if is_torch_available()
			    else {}
			)
			__snake_case  :		str		=      False

			__snake_case  :		str		=      False
			__snake_case  :		List[str]		=      False
			__snake_case  :		int		=      False
			__snake_case  :		Dict		=      False






			def 				_lowercase	(   self  :							Any	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:str										=      DinatModelTester(self	)
					SCREAMING_SNAKE_CASE_						:Optional[Any]										=      ConfigTester(self      ,       config_class=SCREAMING_SNAKE_CASE      ,       embed_dim=37	)






			def 				_lowercase	(   self  :							List[Any]	):




					"""simple docstring"""







					self.create_and_test_config_common_properties()
					self.config_tester.create_and_test_config_to_json_string()
					self.config_tester.create_and_test_config_to_json_file()
					self.config_tester.create_and_test_config_from_and_save_pretrained()
					self.config_tester.create_and_test_config_with_num_labels()
					self.config_tester.check_config_can_be_init_without_params()
					self.config_tester.check_config_arguments_init()






			def 				_lowercase	(   self  :							Union[str, Any]	):




					"""simple docstring"""







					return






			def 				_lowercase	(   self  :							Tuple	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE	)






			def 				_lowercase	(   self  :							int	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:str										=      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE	)






			def 				_lowercase	(   self  :							List[str]	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:int										=      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_backbone(*SCREAMING_SNAKE_CASE	)






			@unittest.skip(reason='Dinat does not use inputs_embeds'	)
			def 				_lowercase	(   self  :							str	):




					"""simple docstring"""







					pass






			@unittest.skip(reason='Dinat does not use feedforward chunking'	)
			def 				_lowercase	(   self  :							Dict	):




					"""simple docstring"""







					pass






			def 				_lowercase	(   self  :							int	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						:int										=      self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							SCREAMING_SNAKE_CASE_						:List[str]										=      model_class(SCREAMING_SNAKE_CASE	)
							self.assertIsInstance(model.get_input_embeddings()      ,       (nn.Module)	)
							SCREAMING_SNAKE_CASE_						:Any										=      model.get_output_embeddings()
							self.assertTrue(x is None or isinstance(SCREAMING_SNAKE_CASE      ,       nn.Linear	)	)






			def 				_lowercase	(   self  :							Tuple	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						:Optional[int]										=      self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							SCREAMING_SNAKE_CASE_						:Tuple										=      model_class(SCREAMING_SNAKE_CASE	)
							SCREAMING_SNAKE_CASE_						:Dict										=      inspect.signature(model.forward	)
							# signature.parameters is an OrderedDict => so arg_names order is deterministic
							SCREAMING_SNAKE_CASE_						:Optional[int]										=      [*signature.parameters.keys()]

							SCREAMING_SNAKE_CASE_						:int										=      ['pixel_values']
							self.assertListEqual(arg_names[:1]      ,       SCREAMING_SNAKE_CASE	)






			def 				_lowercase	(   self  :							Tuple	):




					"""simple docstring"""







					self.skipTest('Dinat\'s attention operation is handled entirely by NATTEN.'	)






			def 				_lowercase	(   self  :							int      ,       SCREAMING_SNAKE_CASE  :							Optional[int]      ,       SCREAMING_SNAKE_CASE  :							Tuple      ,       SCREAMING_SNAKE_CASE  :							List[Any]      ,       SCREAMING_SNAKE_CASE  :							Tuple	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      model_class(SCREAMING_SNAKE_CASE	)
					model.to(SCREAMING_SNAKE_CASE	)
					model.eval()

					with torch.no_grad():
							SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      model(**self._prepare_for_class(SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE	)	)

					SCREAMING_SNAKE_CASE_						:int										=      outputs.hidden_states

					SCREAMING_SNAKE_CASE_						:List[Any]										=      getattr(
					    self.model_tester      ,       'expected_num_hidden_layers'      ,       len(self.model_tester.depths	) + 1	)
					self.assertEqual(len(SCREAMING_SNAKE_CASE	)      ,       SCREAMING_SNAKE_CASE	)

					# Dinat has a different seq_length
					SCREAMING_SNAKE_CASE_						:Optional[Any]										=      (
					    config.patch_size
					    if isinstance(config.patch_size      ,       collections.abc.Iterable	)
					    else (config.patch_size, config.patch_size)
					)

					SCREAMING_SNAKE_CASE_						:List[Any]										=      image_size[0] // patch_size[0]
					SCREAMING_SNAKE_CASE_						:str										=      image_size[1] // patch_size[1]

					self.assertListEqual(
					    list(hidden_states[0].shape[-3:]	)      ,       [height, width, self.model_tester.embed_dim]      ,       )

					if model_class.__name__ != "DinatBackbone":
							SCREAMING_SNAKE_CASE_						:int										=      outputs.reshaped_hidden_states
							self.assertEqual(len(SCREAMING_SNAKE_CASE	)      ,       SCREAMING_SNAKE_CASE	)

							SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						:Dict										=      reshaped_hidden_states[0].shape
							SCREAMING_SNAKE_CASE_						:List[str]										=      (
							    reshaped_hidden_states[0].view(SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE	).permute(0      ,       2      ,       3      ,       1	)
							)
							self.assertListEqual(
							    list(reshaped_hidden_states.shape[-3:]	)      ,       [height, width, self.model_tester.embed_dim]      ,       )






			def 				_lowercase	(   self  :							Any	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						:List[Any]										=      self.model_tester.prepare_config_and_inputs_for_common()

					SCREAMING_SNAKE_CASE_						:str										=      (
					    self.model_tester.image_size
					    if isinstance(self.model_tester.image_size      ,       collections.abc.Iterable	)
					    else (self.model_tester.image_size, self.model_tester.image_size)
					)

					for model_class in self.all_model_classes:
							SCREAMING_SNAKE_CASE_						:str										=      True
							self.check_hidden_states_output(SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE	)

							# check that output_hidden_states also work using config
							del inputs_dict["output_hidden_states"]
							SCREAMING_SNAKE_CASE_						:Tuple										=      True

							self.check_hidden_states_output(SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE      ,       SCREAMING_SNAKE_CASE	)






			@slow
			def 				_lowercase	(   self  :							Optional[int]	):




					"""simple docstring"""







					for model_name in DINAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							SCREAMING_SNAKE_CASE_						:str										=      DinatModel.from_pretrained(SCREAMING_SNAKE_CASE	)
							self.assertIsNotNone(SCREAMING_SNAKE_CASE	)






			def 				_lowercase	(   self  :							Union[str, Any]	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_						:int										=      self.model_tester.prepare_config_and_inputs_for_common()

					SCREAMING_SNAKE_CASE_						:int										=      _config_zero_init(SCREAMING_SNAKE_CASE	)
					for model_class in self.all_model_classes:
							SCREAMING_SNAKE_CASE_						:List[Any]										=      model_class(config=SCREAMING_SNAKE_CASE	)
							for name, param in model.named_parameters():
									if "embeddings" not in name and param.requires_grad:
											self.assertIn(
											    ((param.data.mean() * 1E9).round() / 1E9).item()      ,       [0.0, 1.0]      ,       msg=f'Parameter {name} of model {model_class} seems not properly initialized'      ,       )




@require_natten
@require_vision
@require_torch
class 				__lowerCAmelCase(							unittest.TestCase      ):






			@cached_property
			def 				_lowercase	(   self  :							int	):




					"""simple docstring"""







					return AutoImageProcessor.from_pretrained('shi-labs/dinat-mini-in1k-224'	) if is_vision_available() else None






			@slow
			def 				_lowercase	(   self  :							Any	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Optional[int]										=      DinatForImageClassification.from_pretrained('shi-labs/dinat-mini-in1k-224'	).to(SCREAMING_SNAKE_CASE	)
					SCREAMING_SNAKE_CASE_						:List[Any]										=      self.default_image_processor

					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'	)
					SCREAMING_SNAKE_CASE_						:Tuple										=      image_processor(images=SCREAMING_SNAKE_CASE      ,       return_tensors='pt'	).to(SCREAMING_SNAKE_CASE	)

					# forward pass
					with torch.no_grad():
							SCREAMING_SNAKE_CASE_						:Dict										=      model(**SCREAMING_SNAKE_CASE	)

					# verify the logits
					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      torch.Size((1, 1_000)	)
					self.assertEqual(outputs.logits.shape      ,       SCREAMING_SNAKE_CASE	)
					SCREAMING_SNAKE_CASE_						:Optional[int]										=      torch.tensor([-0.15_45, -0.76_67, 0.46_42]	).to(SCREAMING_SNAKE_CASE	)
					self.assertTrue(torch.allclose(outputs.logits[0, :3]      ,       SCREAMING_SNAKE_CASE      ,       atol=1E-4	)	)




@require_torch
@require_natten
class 				__lowerCAmelCase(							unittest.TestCase		,			lowerCAmelCase__      ):
			__snake_case  :		Optional[int]		=      (DinatBackbone,) if is_torch_available() else ()
			__snake_case  :		Union[str, Any]		=      DinatConfig






			def 				_lowercase	(   self  :							Any	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:List[Any]										=      DinatModelTester(self	)





'''simple docstring'''






import collections

from .utils import ExplicitEnum, is_torch_available, logging


if is_torch_available():
						import torch


lowerCamelCase__	  =		logging.get_logger(__name__)

class 	_lowerCAmelCase     :






		'''simple docstring'''
		def __init__(						self							:       Optional[int]      ,		UpperCamelCase_							:       Optional[Any]      ,		UpperCamelCase_							:       Union[str, Any]=21      ,		UpperCamelCase_							:       Dict=[]      ,		UpperCamelCase_							:       Any=None		)      ->			Any:







									'''simple docstring'''





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

									self.analyse_model()

									self.register_forward_hook()
		def 						__lowercase      (						self							:       Union[str, Any]      ,		UpperCamelCase_							:       Optional[Any]=None		)      ->			Optional[Any]:







									'''simple docstring'''





									if frame is not None:
																self.expand_frame(UpperCamelCase_		)
									self.frames.append('''\n'''.join(self.frame		)		)
									_lowercase							:		Dict				=   []  # start a new frame
		def 						__lowercase      (						self							:       Union[str, Any]      ,		UpperCamelCase_							:       Optional[Any]		)      ->			Union[str, Any]:







									'''simple docstring'''





									self.frame.append(UpperCamelCase_		)
		def 						__lowercase      (						self							:       str		)      ->			List[str]:







									'''simple docstring'''





									print('''\n'''.join(self.frames		)		)
									_lowercase							:		Union[str, Any]				=   []
		def 						__lowercase      (						self							:       Optional[int]		)      ->			List[str]:







									'''simple docstring'''





									_lowercase							:		Optional[Any]				=   []
		def 						__lowercase      (						self							:       List[Any]		)      ->			Optional[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'''		)
									_lowercase							:		List[Any]				=   []
		def 						__lowercase      (						self							:       List[str]		)      ->			Tuple:







									'''simple docstring'''





									_lowercase							:		int				=   {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 						__lowercase      (						self							:       Dict      ,		UpperCamelCase_							:       str      ,		UpperCamelCase_							:       Optional[int]		)      ->			int:







									'''simple docstring'''





									if torch.is_tensor(UpperCamelCase_		):
																self.expand_frame(get_abs_min_max(UpperCamelCase_      ,		UpperCamelCase_		)		)
																if detect_overflow(UpperCamelCase_      ,		UpperCamelCase_		):
																							_lowercase							:		Any				=   True
									elif var is None:
																self.expand_frame(F"{'None':>17} {ctx}"		)
									else:
																self.expand_frame(F"{'not a tensor':>17} {ctx}"		)
		def 						__lowercase      (						self							:       Union[str, Any]		)      ->			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 						__lowercase      (						self							:       str		)      ->			Any:







									'''simple docstring'''





									self.expand_frame(F"{self.prefix} *** Finished batch number={self.batch_number-1} ***\n\n"		)
		def 						__lowercase      (						self							:       Tuple      ,		UpperCamelCase_							:       Any      ,		UpperCamelCase_							:       Union[str, Any]      ,		UpperCamelCase_							:       Optional[int]		)      ->			Dict:







									'''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 						__lowercase      (						self							:       Optional[int]		)      ->			Dict:







									'''simple docstring'''





									self.model.apply(self._register_forward_hook		)
		def 						__lowercase      (						self							:       Optional[int]      ,		UpperCamelCase_							:       int		)      ->			List[str]:







									'''simple docstring'''





									module.register_forward_hook(self.forward_hook		)





		def 						__lowercase      (						self							:       Optional[int]      ,		UpperCamelCase_							:       Optional[int]      ,		UpperCamelCase_							:       Dict      ,		UpperCamelCase_							:       List[Any]		)      ->			List[Any]:







									'''simple docstring'''





									_lowercase							:		str				=   False

									_lowercase							:		Dict				=   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
																_lowercase							:		Union[str, Any]				=   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 				_SCREAMING_SNAKE_CASE(					snake_case_    :     Any ,       snake_case_    :     Dict  )		->Any:






							'''simple docstring'''
							_lowercase							:		Union[str, Any]				=   var.abs()
							return F"{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}"





def 				_SCREAMING_SNAKE_CASE(					snake_case_    :     int ,       snake_case_    :     Any  )		->int:






							'''simple docstring'''
							_lowercase							:		List[Any]				=   False
							if torch.isnan(snake_case_  ).any().item():
														_lowercase							:		Any				=   True
														print(F"{ctx} has nans"  )
							if torch.isinf(snake_case_  ).any().item():
														_lowercase							:		Optional[Any]				=   True
														print(F"{ctx} has infs"  )

							# if needed to monitor large elements can enable the following
							if 0:  # and detected:
														_lowercase							:		List[Any]				=   var[torch.ge(var.abs() ,       1_00  )]
														if naaa.numel() > 0:
																					print(F"{ctx}:  n100={naaa.numel()}"  )
														_lowercase							:		List[str]				=   var[torch.ge(var.abs() ,       10_00  )]
														if naaaa.numel() > 0:
																					print(F"{ctx}: n1000={naaaa.numel()}"  )
														_lowercase							:		Tuple				=   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 	_lowerCAmelCase     (    __A       ):






		'''simple docstring'''
		snake_case_							= 'underflow_overflow'
		snake_case_							= 'tpu_metrics_debug'



