




import inspect
import unittest

from transformers import BitConfig
from transformers.testing_utils import 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, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch
							from torch import nn

							from transformers import BitBackbone, BitForImageClassification, BitImageProcessor, BitModel
							from transformers.models.bit.modeling_bit import BIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
							from PIL import Image




class 		__A      :

						def __init__(		self   :List[Any]    , __snake_case   :List[str]    , __snake_case   :Dict=3    , __snake_case   :Optional[Any]=32    , __snake_case   :List[str]=3    , __snake_case   :str=10    , __snake_case   :Tuple=[8, 16, 32, 64]    , __snake_case   :Any=[1, 1, 2, 1]    , __snake_case   :List[str]=True    , __snake_case   :int=True    , __snake_case   :List[str]="relu"    , __snake_case   :Tuple=3    , __snake_case   :Dict=None    , __snake_case   :Any=["stage2", "stage3", "stage4"]    , __snake_case   :Tuple=[2, 3, 4]    , __snake_case   :Any=1    , ):
										'''simple docstring'''

										__magic_name__	:	int								=parent
										__magic_name__	:	List[str]								=batch_size
										__magic_name__	:	Tuple								=image_size
										__magic_name__	:	Tuple								=num_channels
										__magic_name__	:	Any								=embeddings_size
										__magic_name__	:	Dict								=hidden_sizes
										__magic_name__	:	Optional[int]								=depths
										__magic_name__	:	Union[str, Any]								=is_training
										__magic_name__	:	Optional[int]								=use_labels
										__magic_name__	:	Optional[Any]								=hidden_act
										__magic_name__	:	Tuple								=num_labels
										__magic_name__	:	Any								=scope
										__magic_name__	:	Optional[int]								=len(__snake_case )
										__magic_name__	:	Optional[int]								=out_features
										__magic_name__	:	List[str]								=out_indices
										__magic_name__	:	Optional[int]								=num_groups

						def     A__						(		self   :Optional[Any] ):
										'''simple docstring'''

										__magic_name__	:	str								=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

										__magic_name__	:	Dict								=None
										if self.use_labels:
														__magic_name__	:	int								=ids_tensor([self.batch_size]    , self.num_labels )

										__magic_name__	:	Optional[Any]								=self.get_config()

										return config, pixel_values, labels

						def     A__						(		self   :str ):
										'''simple docstring'''

										return BitConfig(
										    num_channels=self.num_channels    , embeddings_size=self.embeddings_size    , hidden_sizes=self.hidden_sizes    , depths=self.depths    , hidden_act=self.hidden_act    , num_labels=self.num_labels    , out_features=self.out_features    , out_indices=self.out_indices    , num_groups=self.num_groups    , )

						def     A__						(		self   :List[str]    , __snake_case   :str    , __snake_case   :str    , __snake_case   :List[str] ):
										'''simple docstring'''

										__magic_name__	:	Union[str, Any]								=BitModel(config=__snake_case )
										model.to(__snake_case )
										model.eval()
										__magic_name__	:	Tuple								=model(__snake_case )
										self.parent.assertEqual(
										    result.last_hidden_state.shape    , (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)    , )

						def     A__						(		self   :Dict    , __snake_case   :Optional[int]    , __snake_case   :Any    , __snake_case   :Tuple ):
										'''simple docstring'''

										__magic_name__	:	Optional[Any]								=self.num_labels
										__magic_name__	:	Tuple								=BitForImageClassification(__snake_case )
										model.to(__snake_case )
										model.eval()
										__magic_name__	:	Union[str, Any]								=model(__snake_case    , labels=__snake_case )
										self.parent.assertEqual(result.logits.shape    , (self.batch_size, self.num_labels) )

						def     A__						(		self   :List[str]    , __snake_case   :Optional[Any]    , __snake_case   :Tuple    , __snake_case   :Any ):
										'''simple docstring'''

										__magic_name__	:	int								=BitBackbone(config=__snake_case )
										model.to(__snake_case )
										model.eval()
										__magic_name__	:	int								=model(__snake_case )

										# verify feature maps
										self.parent.assertEqual(len(result.feature_maps )    , len(config.out_features ) )
										self.parent.assertListEqual(list(result.feature_maps[0].shape )    , [self.batch_size, self.hidden_sizes[1], 4, 4] )

										# verify channels
										self.parent.assertEqual(len(model.channels )    , len(config.out_features ) )
										self.parent.assertListEqual(model.channels    , config.hidden_sizes[1:] )

										# verify backbone works with out_features=None
										__magic_name__	:	Tuple								=None
										__magic_name__	:	Union[str, Any]								=BitBackbone(config=__snake_case )
										model.to(__snake_case )
										model.eval()
										__magic_name__	:	int								=model(__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, self.hidden_sizes[-1], 1, 1] )

										# verify channels
										self.parent.assertEqual(len(model.channels )    , 1 )
										self.parent.assertListEqual(model.channels    , [config.hidden_sizes[-1]] )


						def     A__						(		self   :str ):
										'''simple docstring'''

										__magic_name__	:	Dict								=self.prepare_config_and_inputs()
										__magic_name__	:	str								=config_and_inputs
										__magic_name__	:	str								={"""pixel_values""": pixel_values}
										return config, inputs_dict




@require_torch
class 		__A      (			UpperCamelCase__					,						UpperCamelCase__					,						unittest.TestCase			):
						UpperCamelCase     						=					(BitModel, BitForImageClassification, BitBackbone) if is_torch_available() else ()
						UpperCamelCase     						=					(
						    {"""feature-extraction""": BitModel, """image-classification""": BitForImageClassification}
						    if is_torch_available()
						    else {}
						)

						UpperCamelCase     						=					False
						UpperCamelCase     						=					False
						UpperCamelCase     						=					False
						UpperCamelCase     						=					False
						UpperCamelCase     						=					False

						def     A__						(		self   :Tuple ):
										'''simple docstring'''

										__magic_name__	:	List[str]								=BitModelTester(self )
										__magic_name__	:	List[Any]								=ConfigTester(self    , config_class=__snake_case    , has_text_modality=__snake_case )

						def     A__						(		self   :Optional[int] ):
										'''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     A__						(		self   :List[Any] ):
										'''simple docstring'''

										return

						@unittest.skip(reason="""Bit does not output attentions""" )
						def     A__						(		self   :Dict ):
										'''simple docstring'''

										pass

						@unittest.skip(reason="""Bit does not use inputs_embeds""" )
						def     A__						(		self   :int ):
										'''simple docstring'''

										pass

						@unittest.skip(reason="""Bit does not support input and output embeddings""" )
						def     A__						(		self   :Dict ):
										'''simple docstring'''

										pass

						def     A__						(		self   :int ):
										'''simple docstring'''

										__magic_name__	:	Dict								=self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														__magic_name__	:	List[str]								=model_class(__snake_case )
														__magic_name__	:	Dict								=inspect.signature(model.forward )
														# signature.parameters is an OrderedDict => so arg_names order is deterministic
														__magic_name__	:	str								=[*signature.parameters.keys()]

														__magic_name__	:	Dict								=["""pixel_values"""]
														self.assertListEqual(arg_names[:1]    , __snake_case )

						def     A__						(		self   :Tuple ):
										'''simple docstring'''

										__magic_name__	:	Tuple								=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*__snake_case )

						def     A__						(		self   :Dict ):
										'''simple docstring'''

										__magic_name__	:	Dict								=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_backbone(*__snake_case )

						def     A__						(		self   :str ):
										'''simple docstring'''

										__magic_name__	:	List[str]								=self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														__magic_name__	:	List[Any]								=model_class(config=__snake_case )
														for name, module in model.named_modules():
																		if isinstance(__snake_case    , (nn.BatchNormad, nn.GroupNorm) ):
																						self.assertTrue(
																						    torch.all(module.weight == 1 )    , msg=f"Parameter {name} of model {model_class} seems not properly initialized"    , )
																						self.assertTrue(
																						    torch.all(module.bias == 0 )    , msg=f"Parameter {name} of model {model_class} seems not properly initialized"    , )

						def     A__						(		self   :str ):
										'''simple docstring'''

										def check_hidden_states_output(__snake_case   :List[str]    , __snake_case   :List[Any]    , __snake_case   :Union[str, Any] ):
														__magic_name__	:	str								=model_class(__snake_case )
														model.to(__snake_case )
														model.eval()

														with torch.no_grad():
																		__magic_name__	:	Tuple								=model(**self._prepare_for_class(__snake_case    , __snake_case ) )

														__magic_name__	:	Dict								=outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

														__magic_name__	:	str								=self.model_tester.num_stages
														self.assertEqual(len(__snake_case )    , expected_num_stages + 1 )

														# Bit's feature maps are of shape (batch_size, num_channels, height, width)
														self.assertListEqual(
														    list(hidden_states[0].shape[-2:] )    , [self.model_tester.image_size // 4, self.model_tester.image_size // 4]    , )

										__magic_name__	:	int								=self.model_tester.prepare_config_and_inputs_for_common()
										__magic_name__	:	Dict								=["""preactivation""", """bottleneck"""]
										for model_class in self.all_model_classes:
														for layer_type in layers_type:
																		__magic_name__	:	Tuple								=layer_type
																		__magic_name__	:	Dict								=True
																		check_hidden_states_output(__snake_case    , __snake_case    , __snake_case )

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

																		check_hidden_states_output(__snake_case    , __snake_case    , __snake_case )

						@unittest.skip(reason="""Bit does not use feedforward chunking""" )
						def     A__						(		self   :Optional[int] ):
										'''simple docstring'''

										pass

						def     A__						(		self   :Tuple ):
										'''simple docstring'''

										__magic_name__	:	Union[str, Any]								=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*__snake_case )


						@slow
						def     A__						(		self   :List[str] ):
										'''simple docstring'''

										for model_name in BIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														__magic_name__	:	Any								=BitModel.from_pretrained(__snake_case )
														self.assertIsNotNone(__snake_case )




def 		lowerCAmelCase_   (	):
				__magic_name__	:	Dict								=Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""   )
				return image




@require_torch
@require_vision
class 		__A      (			unittest.TestCase			):

						@cached_property
						def     A__						(		self   :List[Any] ):
										'''simple docstring'''

										return (
										    BitImageProcessor.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0] ) if is_vision_available() else None
										)


						@slow
						def     A__						(		self   :Union[str, Any] ):
										'''simple docstring'''

										__magic_name__	:	Any								=BitForImageClassification.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0] ).to(__snake_case )

										__magic_name__	:	int								=self.default_image_processor
										__magic_name__	:	List[str]								=prepare_img()
										__magic_name__	:	List[Any]								=image_processor(images=__snake_case    , return_tensors="""pt""" ).to(__snake_case )

										# forward pass
										with torch.no_grad():
														__magic_name__	:	Optional[Any]								=model(**__snake_case )

										# verify the logits
										__magic_name__	:	Tuple								=torch.Size((1, 10_00) )
										self.assertEqual(outputs.logits.shape    , __snake_case )

										__magic_name__	:	Dict								=torch.tensor([[-0.6526, -0.5263, -1.4398]] ).to(__snake_case )

										self.assertTrue(torch.allclose(outputs.logits[0, :3]    , __snake_case    , atol=1E-4 ) )




@require_torch
class 		__A      (			UpperCamelCase__					,						unittest.TestCase			):
						UpperCamelCase     						=					(BitBackbone,) if is_torch_available() else ()
						UpperCamelCase     						=					BitConfig

						UpperCamelCase     						=					False

						def     A__						(		self   :Optional[int] ):
										'''simple docstring'''

										__magic_name__	:	Tuple								=BitModelTester(self )






import random
import timeit
from functools import wraps
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_tf_auto import TF_MODEL_MAPPING, TF_MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_tf_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_tf_available():
							import tensorflow as tf
							from tensorflow.python.framework.errors_impl import ResourceExhaustedError

							from .benchmark_args_tf import TensorFlowBenchmarkArguments

if is_pyanvml_available():
							import pyanvml.pyanvml as nvml

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




def 		lowerCAmelCase_   (	lowerCamelCase	,     lowerCamelCase   ):
				def run_func(lowerCamelCase   ):
								@wraps(lowerCamelCase   )
								def run_in_eager_mode(*lowerCamelCase	,     **lowerCamelCase   ):
												return func(*lowerCamelCase	,     **lowerCamelCase   )

								@wraps(lowerCamelCase   )
								@tf.function(experimental_compile=lowerCamelCase   )
								def run_in_graph_mode(*lowerCamelCase	,     **lowerCamelCase   ):
												return func(*lowerCamelCase	,     **lowerCamelCase   )

								if do_eager_mode is True:
												if use_xla is not False:
																raise ValueError(
																    """Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`."""   )
												return run_in_eager_mode
								else:
												return run_in_graph_mode

				return run_func




def 		lowerCAmelCase_   (	lowerCamelCase	,     lowerCamelCase	,     lowerCamelCase   ):
				__magic_name__	:	Dict								=random.Random()
				__magic_name__	:	Union[str, Any]								=[rng.randint(0	,     vocab_size - 1   ) for i in range(batch_size * sequence_length   )]
				return tf.constant(lowerCamelCase	,     shape=(batch_size, sequence_length)	,     dtype=tf.intaa   )




class 		__A      (			UpperCamelCase__			):
						UpperCamelCase     						=					42
						UpperCamelCase     						=					42
						UpperCamelCase     						=					"TensorFlow"

						@property
						def     A__						(		self   :str ):
										'''simple docstring'''

										return tf.__version__

						def     A__						(		self   :str    , __snake_case   :str    , __snake_case   :int    , __snake_case   :int ):
										'''simple docstring'''

										__magic_name__	:	Dict								=self.args.strategy
										if strategy is None:
														raise ValueError("""A device strategy has to be initialized before using TensorFlow.""" )
										__magic_name__	:	Union[str, Any]								=self._prepare_inference_func(__snake_case    , __snake_case    , __snake_case )
										return self._measure_speed(_inference )

						def     A__						(		self   :int    , __snake_case   :str    , __snake_case   :int    , __snake_case   :int ):
										'''simple docstring'''

										__magic_name__	:	Tuple								=self.args.strategy
										if strategy is None:
														raise ValueError("""A device strategy has to be initialized before using TensorFlow.""" )
										__magic_name__	:	Any								=self._prepare_train_func(__snake_case    , __snake_case    , __snake_case )
										return self._measure_speed(_train )

						def     A__						(		self   :str    , __snake_case   :str    , __snake_case   :int    , __snake_case   :int ):
										'''simple docstring'''

										if self.args.is_gpu:
														tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]    , __snake_case )
										__magic_name__	:	int								=self.args.strategy
										if strategy is None:
														raise ValueError("""A device strategy has to be initialized before using TensorFlow.""" )
										__magic_name__	:	Tuple								=self._prepare_inference_func(__snake_case    , __snake_case    , __snake_case )
										return self._measure_memory(_inference )

						def     A__						(		self   :str    , __snake_case   :str    , __snake_case   :int    , __snake_case   :int ):
										'''simple docstring'''

										if self.args.is_gpu:
														tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]    , __snake_case )
										__magic_name__	:	Tuple								=self.args.strategy
										if strategy is None:
														raise ValueError("""A device strategy has to be initialized before using TensorFlow.""" )

										__magic_name__	:	Any								=self._prepare_train_func(__snake_case    , __snake_case    , __snake_case )
										return self._measure_memory(_train )

						def     A__						(		self   :int    , __snake_case   :str    , __snake_case   :int    , __snake_case   :int ):
										'''simple docstring'''

										__magic_name__	:	int								=self.config_dict[model_name]

										if self.args.fpaa:
														raise NotImplementedError("""Mixed precision is currently not supported.""" )

										__magic_name__	:	Any								=(
										    hasattr(__snake_case    , """architectures""" )
										    and isinstance(config.architectures    , __snake_case )
										    and len(config.architectures ) > 0
										)
										if not self.args.only_pretrain_model and has_model_class_in_config:
														try:
																		__magic_name__	:	Optional[int]								="""TF""" + config.architectures[0]  # prepend 'TF' for tensorflow model
																		__magic_name__	:	Optional[int]								=__import__("""transformers"""    , fromlist=[model_class] )
																		__magic_name__	:	Optional[Any]								=getattr(__snake_case    , __snake_case )
																		__magic_name__	:	Optional[Any]								=model_cls(__snake_case )
														except ImportError:
																		raise ImportError(
																		    f"{model_class} does not exist. If you just want to test the pretrained model, you might want to"
																		    """ set `--only_pretrain_model` or `args.only_pretrain_model=True`.""" )
										else:
														__magic_name__	:	Optional[int]								=TF_MODEL_MAPPING[config.__class__](__snake_case )

										# encoder-decoder has vocab size saved differently
										__magic_name__	:	List[str]								=config.vocab_size if hasattr(__snake_case    , """vocab_size""" ) else config.encoder.vocab_size
										__magic_name__	:	Any								=random_input_ids(__snake_case    , __snake_case    , __snake_case )

										@run_with_tf_optimizations(self.args.eager_mode    , self.args.use_xla )
										def encoder_decoder_forward():
														return model(__snake_case    , decoder_input_ids=__snake_case    , training=__snake_case )

										@run_with_tf_optimizations(self.args.eager_mode    , self.args.use_xla )
										def encoder_forward():
														return model(__snake_case    , training=__snake_case )

										__magic_name__	:	Tuple								=encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

										return _inference

						def     A__						(		self   :int    , __snake_case   :str    , __snake_case   :int    , __snake_case   :int ):
										'''simple docstring'''

										__magic_name__	:	Tuple								=self.config_dict[model_name]

										if self.args.eager_mode is not False:
														raise ValueError("""Training cannot be done in eager mode. Please make sure that `args.eager_mode = False`.""" )

										if self.args.fpaa:
														raise NotImplementedError("""Mixed precision is currently not supported.""" )

										__magic_name__	:	int								=(
										    hasattr(__snake_case    , """architectures""" )
										    and isinstance(config.architectures    , __snake_case )
										    and len(config.architectures ) > 0
										)
										if not self.args.only_pretrain_model and has_model_class_in_config:
														try:
																		__magic_name__	:	List[Any]								="""TF""" + config.architectures[0]  # prepend 'TF' for tensorflow model
																		__magic_name__	:	Optional[int]								=__import__("""transformers"""    , fromlist=[model_class] )
																		__magic_name__	:	str								=getattr(__snake_case    , __snake_case )
																		__magic_name__	:	List[Any]								=model_cls(__snake_case )
														except ImportError:
																		raise ImportError(
																		    f"{model_class} does not exist. If you just want to test the pretrained model, you might want to"
																		    """ set `--only_pretrain_model` or `args.only_pretrain_model=True`.""" )
										else:
														__magic_name__	:	int								=TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](__snake_case )

										# encoder-decoder has vocab size saved differently
										__magic_name__	:	int								=config.vocab_size if hasattr(__snake_case    , """vocab_size""" ) else config.encoder.vocab_size
										__magic_name__	:	Optional[Any]								=random_input_ids(__snake_case    , __snake_case    , __snake_case )

										@run_with_tf_optimizations(self.args.eager_mode    , self.args.use_xla )
										def encoder_decoder_train():
														__magic_name__	:	List[str]								=model(__snake_case    , decoder_input_ids=__snake_case    , labels=__snake_case    , training=__snake_case )[0]
														__magic_name__	:	int								=tf.gradients(__snake_case    , model.trainable_variables )
														return gradients

										@run_with_tf_optimizations(self.args.eager_mode    , self.args.use_xla )
										def encoder_train():
														__magic_name__	:	str								=model(__snake_case    , labels=__snake_case    , training=__snake_case )[0]
														__magic_name__	:	int								=tf.gradients(__snake_case    , model.trainable_variables )
														return gradients

										__magic_name__	:	Union[str, Any]								=encoder_decoder_train if config.is_encoder_decoder else encoder_train

										return _train

						def     A__						(		self   :Any    , __snake_case   :Union[str, Any] ):
										'''simple docstring'''

										with self.args.strategy.scope():
														try:
																		if self.args.is_tpu or self.args.use_xla:
																						# run additional 10 times to stabilize compilation for tpu
																						logger.info("""Do inference on TPU. Running model 5 times to stabilize compilation""" )
																						timeit.repeat(__snake_case    , repeat=1    , number=5 )

																		# as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
																		__magic_name__	:	Union[str, Any]								=timeit.repeat(
																		    __snake_case    , repeat=self.args.repeat    , number=10    , )

																		return min(__snake_case ) / 10.0
														except ResourceExhaustedError as e:
																		self.print_fn(f"Doesn't fit on GPU. {e}" )


						def     A__						(		self   :Any    , __snake_case   :Callable[[], None] ):
										'''simple docstring'''

										logger.info(
										    """Note that TensorFlow allocates more memory than """
										    """it might need to speed up computation. """
										    """The memory reported here corresponds to the memory """
										    """reported by `nvidia-smi`, which can vary depending """
										    """on total available memory on the GPU that is used.""" )
										with self.args.strategy.scope():
														try:
																		if self.args.trace_memory_line_by_line:
																						if not self.args.eager_mode:
																										raise ValueError(
																										    """`args.eager_mode` is set to `False`. Make sure to run model in eager mode to measure memory"""
																										    """ consumption line by line.""" )
																						__magic_name__	:	Union[str, Any]								=start_memory_tracing("""transformers""" )

																		if self.args.is_tpu:
																						# tpu
																						raise NotImplementedError(
																						    """Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking"""
																						    """ with `args.memory=False`""" )
																		elif self.args.is_gpu:
																						# gpu
																						if not is_pyanvml_available():
																										logger.warning(
																										    """py3nvml not installed, we won't log GPU memory usage. """
																										    """Install py3nvml (pip install py3nvml) to log information about GPU.""" )
																										__magic_name__	:	str								="""N/A"""
																						else:
																										logger.info(
																										    """Measuring total GPU usage on GPU device. Make sure to not have additional processes"""
																										    """ running on the same GPU.""" )
																										# init nvml
																										nvml.nvmlInit()
																										func()
																										__magic_name__	:	List[str]								=nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx )
																										__magic_name__	:	Tuple								=nvml.nvmlDeviceGetMemoryInfo(__snake_case )
																										__magic_name__	:	Any								=meminfo.used
																										__magic_name__	:	str								=Memory(__snake_case )
																										# shutdown nvml
																										nvml.nvmlShutdown()
																		else:
																						# cpu
																						if self.args.trace_memory_line_by_line:
																										logger.info(
																										    """When enabling line by line tracing, the max peak memory for CPU is inaccurate in"""
																										    """ TensorFlow.""" )
																										__magic_name__	:	List[str]								=None
																						else:
																										__magic_name__	:	List[Any]								=measure_peak_memory_cpu(__snake_case )
																										__magic_name__	:	str								=Memory(__snake_case ) if isinstance(__snake_case    , __snake_case ) else memory_bytes
																		if self.args.trace_memory_line_by_line:
																						__magic_name__	:	List[Any]								=stop_memory_tracing(__snake_case )
																						if memory is None:
																										__magic_name__	:	Any								=summary.total
																		else:
																						__magic_name__	:	Optional[Any]								=None

																		return memory, summary
														except ResourceExhaustedError as e:
																		self.print_fn(f"Doesn't fit on GPU. {e}" )
																		return "N/A", None




