




from abc import ABC, abstractmethod
from typing import List, Optional




class 		__A      (			__a			):

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

										self.test()

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

										__magic_name__	:	Tuple								=0
										__magic_name__	:	List[Any]								=False
										while not completed:
														if counter == 1:
																		self.reset()
														__magic_name__	:	List[str]								=self.advance()
														if not self.does_advance(a_ ):
																		raise Exception(
																		    """Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true.""" )

														__magic_name__	:	Union[str, Any]								=self.update(a_ )
														counter += 1

														if counter > 1_00_00:
																		raise Exception("""update() does not fulfill the constraint.""" )

										if self.remaining() != 0:
														raise Exception("""Custom Constraint is not defined correctly.""" )

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

										raise NotImplementedError(
										    f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." )

						@abstractmethod
						def     A__						(		self   :Optional[int]    , __snake_case   :int ):
										'''simple docstring'''

										raise NotImplementedError(
										    f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." )

						@abstractmethod
						def     A__						(		self   :Optional[int]    , __snake_case   :int ):
										'''simple docstring'''

										raise NotImplementedError(
										    f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." )

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

										raise NotImplementedError(
										    f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." )

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

										raise NotImplementedError(
										    f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." )


						@abstractmethod
						def     A__						(		self   :List[str]    , __snake_case   :Tuple=False ):
										'''simple docstring'''

										raise NotImplementedError(
										    f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." )




class 		__A      (			__a			):

						def __init__(		self   :Tuple    , __snake_case   :List[int] ):
										'''simple docstring'''

										super(a_    , self ).__init__()

										if not isinstance(a_    , a_ ) or len(a_ ) == 0:
														raise ValueError(f"`token_ids` has to be a non-empty list, but is {token_ids}." )
										if any((not isinstance(a_    , a_ ) or token_id < 0) for token_id in token_ids ):
														raise ValueError(f"Each list in `token_ids` has to be a list of positive integers, but is {token_ids}." )

										__magic_name__	:	Optional[int]								=token_ids

										__magic_name__	:	Dict								=len(self.token_ids )
										__magic_name__	:	Union[str, Any]								=-1  # the index of the currently fulfilled step
										__magic_name__	:	str								=False

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

										if self.completed:
														return None
										return self.token_ids[self.fulfilled_idx + 1]

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

										if not isinstance(a_    , a_ ):
														raise ValueError(f"`token_id` has to be an `int`, but is {token_id} of type {type(a_ )}" )

										if self.completed:
														return False

										return token_id == self.token_ids[self.fulfilled_idx + 1]

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

										if not isinstance(a_    , a_ ):
														raise ValueError(f"`token_id` has to be an `int`, but is {token_id} of type {type(a_ )}" )

										__magic_name__	:	Optional[int]								=False
										__magic_name__	:	List[str]								=False
										__magic_name__	:	List[Any]								=False

										if self.does_advance(a_ ):
														self.fulfilled_idx += 1
														__magic_name__	:	List[str]								=True
														if self.fulfilled_idx == (self.seqlen - 1):
																		__magic_name__	:	str								=True
														__magic_name__	:	Optional[int]								=completed
										else:
														# failed to make progress.
														__magic_name__	:	Union[str, Any]								=True
														self.reset()
										return stepped, completed, reset

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

										__magic_name__	:	int								=False
										__magic_name__	:	Dict								=0

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

										return self.seqlen - (self.fulfilled_idx + 1)


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

										__magic_name__	:	Dict								=PhrasalConstraint(self.token_ids )

										if stateful:
														__magic_name__	:	List[Any]								=self.seqlen
														__magic_name__	:	str								=self.fulfilled_idx
														__magic_name__	:	Union[str, Any]								=self.completed

										return new_constraint




class 		__A      :

						def __init__(		self   :int    , __snake_case   :List[List[int]]    , __snake_case   :int=True ):
										'''simple docstring'''

										__magic_name__	:	List[str]								=max([len(a_ ) for one in nested_token_ids] )

										__magic_name__	:	int								={}
										for token_ids in nested_token_ids:
														__magic_name__	:	List[str]								=root
														for tidx, token_id in enumerate(a_ ):
																		if token_id not in level:
																						__magic_name__	:	int								={}

																		__magic_name__	:	List[Any]								=level[token_id]

										if no_subsets and self.has_subsets(a_    , a_ ):
														raise ValueError(
														    """Each list in `nested_token_ids` can't be a complete subset of another list, but is"""
														    f" {nested_token_ids}." )

										__magic_name__	:	int								=root

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

										__magic_name__	:	Union[str, Any]								=self.trie

										for current_token in current_seq:
														__magic_name__	:	Any								=start[current_token]

										__magic_name__	:	Any								=list(start.keys() )

										return next_tokens

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

										__magic_name__	:	Any								=self.next_tokens(a_ )

										return len(a_ ) == 0

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

										__magic_name__	:	int								=list(root.values() )
										if len(a_ ) == 0:
														return 1
										else:
														return sum([self.count_leaves(a_ ) for nn in next_nodes] )


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

										__magic_name__	:	str								=self.count_leaves(a_ )
										return len(a_ ) != leaf_count




class 		__A      (			__a			):

						def __init__(		self   :Tuple    , __snake_case   :List[List[int]] ):
										'''simple docstring'''

										super(a_    , self ).__init__()

										if not isinstance(a_    , a_ ) or len(a_ ) == 0:
														raise ValueError(f"`nested_token_ids` has to be a non-empty list, but is {nested_token_ids}." )
										if any(not isinstance(a_    , a_ ) for token_ids in nested_token_ids ):
														raise ValueError(f"`nested_token_ids` has to be a list of lists, but is {nested_token_ids}." )
										if any(
										    any((not isinstance(a_    , a_ ) or token_id < 0) for token_id in token_ids )
										    for token_ids in nested_token_ids ):
														raise ValueError(
														    f"Each list in `nested_token_ids` has to be a list of positive integers, but is {nested_token_ids}." )

										__magic_name__	:	Optional[Any]								=DisjunctiveTrie(a_ )
										__magic_name__	:	Optional[int]								=nested_token_ids

										__magic_name__	:	Tuple								=self.trie.max_height
										__magic_name__	:	List[str]								=[]
										__magic_name__	:	Union[str, Any]								=False

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

										__magic_name__	:	List[Any]								=self.trie.next_tokens(self.current_seq )

										if len(a_ ) == 0:
														return None
										else:
														return token_list

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

										if not isinstance(a_    , a_ ):
														raise ValueError(f"`token_id` is supposed to be type `int`, but is {token_id} of type {type(a_ )}" )

										__magic_name__	:	Optional[Any]								=self.trie.next_tokens(self.current_seq )

										return token_id in next_tokens

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

										if not isinstance(a_    , a_ ):
														raise ValueError(f"`token_id` is supposed to be type `int`, but is {token_id} of type {type(a_ )}" )

										__magic_name__	:	Any								=False
										__magic_name__	:	Union[str, Any]								=False
										__magic_name__	:	Tuple								=False

										if self.does_advance(a_ ):
														self.current_seq.append(a_ )
														__magic_name__	:	Optional[Any]								=True
										else:
														__magic_name__	:	Tuple								=True
														self.reset()

										__magic_name__	:	Any								=self.trie.reached_leaf(self.current_seq )
										__magic_name__	:	Tuple								=completed

										return stepped, completed, reset

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

										__magic_name__	:	Any								=False
										__magic_name__	:	Tuple								=[]

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

										if self.completed:
														# since this can be completed without reaching max height
														return 0
										else:
														return self.seqlen - len(self.current_seq )


						def     A__						(		self   :Dict    , __snake_case   :List[Any]=False ):
										'''simple docstring'''

										__magic_name__	:	Union[str, Any]								=DisjunctiveConstraint(self.token_ids )

										if stateful:
														__magic_name__	:	Tuple								=self.seqlen
														__magic_name__	:	List[str]								=self.current_seq
														__magic_name__	:	List[Any]								=self.completed

										return new_constraint




class 		__A      :

						def __init__(		self   :Any    , __snake_case   :List[Constraint] ):
										'''simple docstring'''

										__magic_name__	:	Any								=constraints

										# max # of steps required to fulfill a given constraint
										__magic_name__	:	Optional[int]								=max([c.seqlen for c in constraints] )
										__magic_name__	:	Optional[Any]								=len(a_ )
										__magic_name__	:	List[Any]								=False

										self.init_state()

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

										__magic_name__	:	List[Any]								=[]
										__magic_name__	:	Union[str, Any]								=None
										__magic_name__	:	Optional[int]								=[constraint.copy(stateful=a_ ) for constraint in self.constraints]

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

										__magic_name__	:	Any								=0
										if self.inprogress_constraint:
														# extra points for having a constraint mid-fulfilled
														add += self.max_seqlen - self.inprogress_constraint.remaining()

										return (len(self.complete_constraints ) * self.max_seqlen) + add

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

										__magic_name__	:	Union[str, Any]								=[]
										if self.inprogress_constraint is None:
														for constraint in self.pending_constraints:  # "pending" == "unfulfilled yet"
																		__magic_name__	:	str								=constraint.advance()
																		if isinstance(a_    , a_ ):
																						token_list.append(a_ )
																		elif isinstance(a_    , a_ ):
																						token_list.extend(a_ )
										else:
														__magic_name__	:	List[str]								=self.inprogress_constraint.advance()
														if isinstance(a_    , a_ ):
																		token_list.append(a_ )
														elif isinstance(a_    , a_ ):
																		token_list.extend(a_ )

										if len(a_ ) == 0:
														return None
										else:
														return token_list

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

										self.init_state()

										if token_ids is not None:
														for token in token_ids:
																		# completes or steps **one** constraint
																		__magic_name__	:	int								=self.add(a_ )

																		# the entire list of constraints are fulfilled
																		if self.completed:
																						break

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

										if not isinstance(a_    , a_ ):
														raise ValueError(f"`token_id` should be an `int`, but is `{token_id}`." )

										__magic_name__	:	List[str]								=False, False

										if self.completed:
														__magic_name__	:	Dict								=True
														__magic_name__	:	List[Any]								=False
														return complete, stepped

										if self.inprogress_constraint is not None:
														# In the middle of fulfilling a constraint. If the `token_id` *does* makes an incremental progress to current
														# job, simply update the state

														__magic_name__	:	Optional[int]								=self.inprogress_constraint.update(a_ )
														if reset:
																		# 1. If the next token breaks the progress, then we must restart.
																		#     e.g. constraint = "I love pies" and sequence so far is "I love" but `token_id` == "books".

																		#     But that doesn't mean we self.init_state(), since we only reset the state for this particular
																		#     constraint, not the full list of constraints.

																		self.pending_constraints.append(self.inprogress_constraint.copy(stateful=a_ ) )
																		__magic_name__	:	Union[str, Any]								=None

														if complete:
																		# 2. If the next token completes the constraint, move it to completed list, set
																		#     inprogress to None. If there are no pending constraints either, then this full list of constraints
																		#     is complete.

																		self.complete_constraints.append(self.inprogress_constraint )
																		__magic_name__	:	Dict								=None

																		if len(self.pending_constraints ) == 0:
																						# we're done!
																						__magic_name__	:	Union[str, Any]								=True

										else:
														# Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards any of our list
														# of constraints?

														for cidx, pending_constraint in enumerate(self.pending_constraints ):
																		if pending_constraint.does_advance(a_ ):
																						__magic_name__	:	Any								=pending_constraint.update(a_ )

																						if not stepped:
																										raise Exception(
																										    """`constraint.update(token_id)` is not yielding incremental progress, """
																										    """even though `constraint.does_advance(token_id)` is true.""" )

																						if complete:
																										self.complete_constraints.append(a_ )
																										__magic_name__	:	Optional[Any]								=None

																						if not complete and stepped:
																										__magic_name__	:	List[Any]								=pending_constraint

																						if complete or stepped:
																										# If we made any progress at all, then it's at least not a "pending constraint".

																										__magic_name__	:	Any								=(
																										    self.pending_constraints[:cidx] + self.pending_constraints[cidx + 1 :]
																										)

																										if len(self.pending_constraints ) == 0 and self.inprogress_constraint is None:
																														# If there's no longer any pending after this and no inprogress either, then we must be
																														# complete.

																														__magic_name__	:	List[str]								=True

																										break  # prevent accidentally stepping through multiple constraints with just one token.

										return complete, stepped


						def     A__						(		self   :Dict    , __snake_case   :int=True ):
										'''simple docstring'''

										__magic_name__	:	Dict								=ConstraintListState(self.constraints )  # we actually never though self.constraints objects
										# throughout this process. So it's at initialization state.

										if stateful:
														__magic_name__	:	List[Any]								=[
														    constraint.copy(stateful=a_ ) for constraint in self.complete_constraints
														]
														if self.inprogress_constraint is not None:
																		__magic_name__	:	Dict								=self.inprogress_constraint.copy(stateful=a_ )
														__magic_name__	:	List[str]								=[constraint.copy() for constraint in self.pending_constraints]

										return new_state






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




