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

snake_case__				= logging.get_logger(__name__)
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase		)						->    List[str]:



      '''simple docstring'''

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

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

            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       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)						->    ["tf.Tensor"]:



      '''simple docstring'''

      _lowerCamelCase			    =				random.Random()
      _lowerCamelCase			    =				[rng.randint(0 ,						vocab_size - 1		) for i in range(batch_size * sequence_length		)]
      return tf.constant(__UpperCAmelCase ,						shape=(batch_size, sequence_length) ,						dtype=tf.intaa		)




class  UpperCamelCase				( __lowercase ):

     '''simple docstring'''

     A_										=  42
     A_										=  42
     A_										=  'TensorFlow'



     @property
     def 		UpperCamelCase_						(					self						) ->				Optional[int]:






           """simple docstring"""
           return tf.__version__



     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_						) ->				float:






           """simple docstring"""
           _lowerCamelCase			    =				self.args.strategy
           if strategy is None:
                 raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''						)
           _lowerCamelCase			    =				self._prepare_inference_func(A_				,				A_				,				A_						)
           return self._measure_speed(_inference						)



     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_						) ->				float:






           """simple docstring"""
           _lowerCamelCase			    =				self.args.strategy
           if strategy is None:
                 raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''						)
           _lowerCamelCase			    =				self._prepare_train_func(A_				,				A_				,				A_						)
           return self._measure_speed(_train						)



     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_						) ->				[Memory, Optional[MemorySummary]]:






           """simple docstring"""
           # initialize GPU on separate process
           if self.args.is_gpu:
                 tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]				,				A_						)
           _lowerCamelCase			    =				self.args.strategy
           if strategy is None:
                 raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''						)
           _lowerCamelCase			    =				self._prepare_inference_func(A_				,				A_				,				A_						)
           return self._measure_memory(_inference						)



     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_						) ->				[Memory, Optional[MemorySummary]]:






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

           _lowerCamelCase			    =				self._prepare_train_func(A_				,				A_				,				A_						)
           return self._measure_memory(_train						)



     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_						) ->				Callable[[], None]:






           """simple docstring"""
           _lowerCamelCase			    =				self.config_dict[model_name]

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

           _lowerCamelCase			    =				(
               hasattr(A_				,				'''architectures'''						)
               and isinstance(config.architectures				,				A_						)
               and len(config.architectures						) > 0
           )
           if not self.args.only_pretrain_model and has_model_class_in_config:
                 try:
                       _lowerCamelCase			    =				'''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
                       _lowerCamelCase			    =				__import__('''transformers'''				,				fromlist=[model_class]						)
                       _lowerCamelCase			    =				getattr(A_				,				A_						)
                       _lowerCamelCase			    =				model_cls(A_						)
                 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:
                 _lowerCamelCase			    =				TF_MODEL_MAPPING[config.__class__](A_						)

           # encoder-decoder has vocab size saved differently
           _lowerCamelCase			    =				config.vocab_size if hasattr(A_				,				'''vocab_size'''						) else config.encoder.vocab_size
           _lowerCamelCase			    =				random_input_ids(A_				,				A_				,				A_						)

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

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

           _lowerCamelCase			    =				encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

           return _inference



     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_						) ->				Callable[[], None]:






           """simple docstring"""
           _lowerCamelCase			    =				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.'''						)

           _lowerCamelCase			    =				(
               hasattr(A_				,				'''architectures'''						)
               and isinstance(config.architectures				,				A_						)
               and len(config.architectures						) > 0
           )
           if not self.args.only_pretrain_model and has_model_class_in_config:
                 try:
                       _lowerCamelCase			    =				'''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
                       _lowerCamelCase			    =				__import__('''transformers'''				,				fromlist=[model_class]						)
                       _lowerCamelCase			    =				getattr(A_				,				A_						)
                       _lowerCamelCase			    =				model_cls(A_						)
                 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:
                 _lowerCamelCase			    =				TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](A_						)

           # encoder-decoder has vocab size saved differently
           _lowerCamelCase			    =				config.vocab_size if hasattr(A_				,				'''vocab_size'''						) else config.encoder.vocab_size
           _lowerCamelCase			    =				random_input_ids(A_				,				A_				,				A_						)

           @run_with_tf_optimizations(self.args.eager_mode				,				self.args.use_xla						)
           def encoder_decoder_train():
                 _lowerCamelCase			    =				model(A_				,				decoder_input_ids=A_				,				labels=A_				,				training=A_						)[0]
                 _lowerCamelCase			    =				tf.gradients(A_				,				model.trainable_variables						)
                 return gradients

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

           _lowerCamelCase			    =				encoder_decoder_train if config.is_encoder_decoder else encoder_train

           return _train



     def 		UpperCamelCase_						(					self				,				A_						) ->				float:






           """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(A_				,				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
                       _lowerCamelCase			    =				timeit.repeat(
                           A_				,				repeat=self.args.repeat				,				number=10				,				)

                       return min(A_						) / 10.0
                 except ResourceExhaustedError as e:
                       self.print_fn(F'Doesn\'t fit on GPU. {e}'						)





     def 		UpperCamelCase_						(					self				,				A_						) ->				[Memory, MemorySummary]:






           """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.'''						)
                             _lowerCamelCase			    =				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.'''						)
                                   _lowerCamelCase			    =				'''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()
                                   _lowerCamelCase			    =				nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx						)
                                   _lowerCamelCase			    =				nvml.nvmlDeviceGetMemoryInfo(A_						)
                                   _lowerCamelCase			    =				meminfo.used
                                   _lowerCamelCase			    =				Memory(A_						)
                                   # 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.'''						)
                                   _lowerCamelCase			    =				None
                             else:
                                   _lowerCamelCase			    =				measure_peak_memory_cpu(A_						)
                                   _lowerCamelCase			    =				Memory(A_						) if isinstance(A_				,				A_						) else memory_bytes
                       if self.args.trace_memory_line_by_line:
                             _lowerCamelCase			    =				stop_memory_tracing(A_						)
                             if memory is None:
                                   _lowerCamelCase			    =				summary.total
                       else:
                             _lowerCamelCase			    =				None

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


def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase		)						->    bool:



      '''simple docstring'''

      _lowerCamelCase			    =				len(__UpperCAmelCase		)
      _lowerCamelCase			    =				[[False] * (required_sum + 1) for _ in range(arr_len + 1		)]

      # for each arr value, a sum of zero(0) can be formed by not taking any element
      # hence True/1
      for i in range(arr_len + 1		):
            _lowerCamelCase			    =				True

      # sum is not zero and set is empty then false
      for i in range(1 ,						required_sum + 1		):
            _lowerCamelCase			    =				False

      for i in range(1 ,						arr_len + 1		):
            for j in range(1 ,						required_sum + 1		):
                  if arr[i - 1] > j:
                        _lowerCamelCase			    =				subset[i - 1][j]
                  if arr[i - 1] <= j:
                        _lowerCamelCase			    =				subset[i - 1][j] or subset[i - 1][j - arr[i - 1]]

      return subset[arr_len][required_sum]


if __name__ == "__main__":
     import doctest

     doctest.testmod()