



"""simple docstring"""







from math import pow

def     UpperCamelCase_   (   lowerCAmelCase__				:    int   ,       lowerCAmelCase__				:    int   ,       lowerCAmelCase__				:    int   ,       lowerCAmelCase__				:    int   ,       lowerCAmelCase__				:    int   ,       )						->       tuple[int, int]:

  """simple docstring"""




  if current_sum == needed_sum:
    # If the sum of the powers is equal to needed_sum, then we have a solution.
    solutions_count += 1
    return current_sum, solutions_count

  lowerCAmelCase_       :    Tuple								=			int(pow(lowerCAmelCase__   ,       lowerCAmelCase__						)						)
  if current_sum + i_to_n <= needed_sum:
    # If the sum of the powers is less than needed_sum, then continue adding powers.
    current_sum += i_to_n
    lowerCAmelCase_						,lowerCAmelCase_       :    str								=			backtrack(
        lowerCAmelCase__   ,       lowerCAmelCase__   ,       current_number + 1   ,       lowerCAmelCase__   ,       lowerCAmelCase__						)
    current_sum -= i_to_n
  if i_to_n < needed_sum:
    # If the power of i is less than needed_sum, then try with the next power.
    lowerCAmelCase_						,lowerCAmelCase_       :    Optional[int]								=			backtrack(
        lowerCAmelCase__   ,       lowerCAmelCase__   ,       current_number + 1   ,       lowerCAmelCase__   ,       lowerCAmelCase__						)
  return current_sum, solutions_count

def     UpperCamelCase_   (   lowerCAmelCase__				:    int   ,       lowerCAmelCase__				:    int						)						->       int:

  """simple docstring"""




  if not (1 <= needed_sum <= 1000 and 2 <= power <= 10):
    raise ValueError(
        'Invalid input\n'
        'needed_sum must be between 1 and 1000, power between 2 and 10.'						)

  return backtrack(lowerCAmelCase__   ,       lowerCAmelCase__   ,       1   ,       0   ,       0						)[1]  # Return the solutions_count


if __name__ == "__main__":
     import doctest

     doctest.testmod()










"""simple docstring"""







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

lowercase__		:    str				   =     logging.get_logger(__name__)

def     UpperCamelCase_   (   lowerCAmelCase__				:    bool   ,       lowerCAmelCase__				:    bool						)						->       List[Any]:

  """simple docstring"""




  def run_func(lowerCAmelCase__				:    int						):
    @wraps(lowerCAmelCase__						)
    def run_in_eager_mode(*lowerCAmelCase__				:    List[Any]   ,       **lowerCAmelCase__				:    int						):
      return func(*lowerCAmelCase__   ,       **lowerCAmelCase__						)

    @wraps(lowerCAmelCase__						)
    @tf.function(experimental_compile=lowerCAmelCase__						)
    def run_in_graph_mode(*lowerCAmelCase__				:    Dict   ,       **lowerCAmelCase__				:    Any						):
      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     UpperCamelCase_   (   lowerCAmelCase__				:    int   ,       lowerCAmelCase__				:    int   ,       lowerCAmelCase__				:    int						)						->       ["tf.Tensor"]:

  """simple docstring"""




  lowerCAmelCase_       :    Dict								=			random.Random()
  lowerCAmelCase_       :    List[str]								=			[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 				UpperCamelCase__       (     lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE      =					42
      _SCREAMING_SNAKE_CASE      =					42
      _SCREAMING_SNAKE_CASE      =					"TensorFlow"



      @property
      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[Any]	):
        return tf.__version__



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				str   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				int	):
        # initialize GPU on separate process
        lowerCAmelCase_       :    List[Any]								=			self.args.strategy
        if strategy is None:
          raise ValueError('A device strategy has to be initialized before using TensorFlow.'	)
        lowerCAmelCase_       :    List[str]								=			self._prepare_inference_func(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)
        return self._measure_speed(_inference	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Any   ,      SCREAMING_SNAKE_CASE_    :				str   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				int	):
        lowerCAmelCase_       :    Optional[int]								=			self.args.strategy
        if strategy is None:
          raise ValueError('A device strategy has to be initialized before using TensorFlow.'	)
        lowerCAmelCase_       :    Any								=			self._prepare_train_func(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)
        return self._measure_speed(_train	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[str]   ,      SCREAMING_SNAKE_CASE_    :				str   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				int	):
        # initialize GPU on separate process
        if self.args.is_gpu:
          tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]   ,      SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    Any								=			self.args.strategy
        if strategy is None:
          raise ValueError('A device strategy has to be initialized before using TensorFlow.'	)
        lowerCAmelCase_       :    Optional[Any]								=			self._prepare_inference_func(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)
        return self._measure_memory(_inference	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Any   ,      SCREAMING_SNAKE_CASE_    :				str   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				int	):
        if self.args.is_gpu:
          tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]   ,      SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    Any								=			self.args.strategy
        if strategy is None:
          raise ValueError('A device strategy has to be initialized before using TensorFlow.'	)

        lowerCAmelCase_       :    Optional[int]								=			self._prepare_train_func(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)
        return self._measure_memory(_train	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Any   ,      SCREAMING_SNAKE_CASE_    :				str   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				int	):
        lowerCAmelCase_       :    Any								=			self.config_dict[model_name]

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

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

        # encoder-decoder has vocab size saved differently
        lowerCAmelCase_       :    List[Any]								=			config.vocab_size if hasattr(SCREAMING_SNAKE_CASE_   ,      'vocab_size'	) else config.encoder.vocab_size
        lowerCAmelCase_       :    Tuple								=			random_input_ids(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)

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

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

        lowerCAmelCase_       :    List[Any]								=			encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

        return _inference



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				str   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				int	):
        lowerCAmelCase_       :    Union[str, Any]								=			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_       :    Dict								=			(
            hasattr(SCREAMING_SNAKE_CASE_   ,      'architectures'	)
            and isinstance(config.architectures   ,      SCREAMING_SNAKE_CASE_	)
            and len(config.architectures	) > 0
        )
        if not self.args.only_pretrain_model and has_model_class_in_config:
          try:
            lowerCAmelCase_       :    Optional[Any]								=			'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
            lowerCAmelCase_       :    int								=			__import__('transformers'   ,      fromlist=[model_class]	)
            lowerCAmelCase_       :    Any								=			getattr(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)
            lowerCAmelCase_       :    Optional[Any]								=			model_cls(SCREAMING_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:
          lowerCAmelCase_       :    Union[str, Any]								=			TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](SCREAMING_SNAKE_CASE_	)

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

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

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

        lowerCAmelCase_       :    int								=			encoder_decoder_train if config.is_encoder_decoder else encoder_train

        return _train



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[str]   ,      SCREAMING_SNAKE_CASE_    :				Tuple	):
        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(SCREAMING_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
            lowerCAmelCase_       :    Dict								=			timeit.repeat(
                SCREAMING_SNAKE_CASE_   ,      repeat=self.args.repeat   ,      number=1_0   ,      )

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





      def 					SCREAMING_SNAKE_CASE__			(				self    :				int   ,      SCREAMING_SNAKE_CASE_    :				Callable[[], None]	):
        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_       :    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.'	)
                lowerCAmelCase_       :    Tuple								=			'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_       :    int								=			nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx	)
                lowerCAmelCase_       :    Union[str, Any]								=			nvml.nvmlDeviceGetMemoryInfo(SCREAMING_SNAKE_CASE_	)
                lowerCAmelCase_       :    List[Any]								=			meminfo.used
                lowerCAmelCase_       :    int								=			Memory(SCREAMING_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.'	)
                lowerCAmelCase_       :    Optional[int]								=			None
              else:
                lowerCAmelCase_       :    Union[str, Any]								=			measure_peak_memory_cpu(SCREAMING_SNAKE_CASE_	)
                lowerCAmelCase_       :    Any								=			Memory(SCREAMING_SNAKE_CASE_	) if isinstance(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	) else memory_bytes
            if self.args.trace_memory_line_by_line:
              lowerCAmelCase_       :    List[Any]								=			stop_memory_tracing(SCREAMING_SNAKE_CASE_	)
              if memory is None:
                lowerCAmelCase_       :    Union[str, Any]								=			summary.total
            else:
              lowerCAmelCase_       :    List[str]								=			None

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





