




"""simple docstring"""

import collections
import json
import math
import os
import re
import time
from fnmatch import fnmatch
from typing import Dict

import requests
from slack_sdk import WebClient


a         =      WebClient(token=os.environ['''CI_SLACK_BOT_TOKEN'''])






def 						_snake_case					( _snake_case							: int							)   -> Dict:






      '''simple docstring'''


      _A        =       test_results.split(' '							)

      _A        =       0
      _A        =       0

      # When the output is short enough, the output is surrounded by = signs: "== OUTPUT =="
      # When it is too long, those signs are not present.
      _A        =       expressions[-2] if '=' in expressions[-1] else expressions[-1]

      for i, expression in enumerate(_snake_case							):
            if "failed" in expression:
                  failed += int(expressions[i - 1]							)
            if "passed" in expression:
                  success += int(expressions[i - 1]							)

      return failed, success, time_spent






def 						_snake_case					( _snake_case							: Union[str, Any]							)   -> List[str]:






      '''simple docstring'''


      _A        =       {}
      _A        =       None
      _A        =       False
      for line in failures_short_lines.split('\n'							):
            if re.search(R'_ \[doctest\]'    ,  _snake_case							):
                  _A        =       True
                  _A        =       line.split(' '							)[2]
            elif in_error and not line.split(' '							)[0].isdigit():
                  _A        =       line
                  _A        =       False

      return failures





class     lowercase_	:



 '''simple docstring'''

 def __init__(			self	:	Optional[Any]   ,							_UpperCAmelCase	:	str   ,							_UpperCAmelCase	:	Dict		):
       _A        =       title

       _A        =       doc_test_results['time_spent'].split(','		)[0]
       _A        =       doc_test_results['success']
       _A        =       doc_test_results['failures']
       _A        =       self.n_success + self.n_failures

       # Failures and success of the modeling tests
       _A        =       doc_test_results

 @property
 def       lowerCAmelCase_							(			self	:	Tuple		):
       _A        =       [self._time_spent]
       _A        =       0

       for time in time_spent:
             _A        =       time.split(':'		)

             # Time can be formatted as xx:xx:xx, as .xx, or as x.xx if the time spent was less than a minute.
             if len(_UpperCAmelCase		) == 1:
                   _A        =       [0, 0, time_parts[0]]

             _A       ,   _A       ,   _A        =       int(time_parts[0]		), int(time_parts[1]		), float(time_parts[2]		)
             total_secs += hours * 3_600 + minutes * 60 + seconds

       _A       ,   _A       ,   _A        =       total_secs // 3_600, (total_secs % 3_600) // 60, total_secs % 60
       return F'''{int(_UpperCAmelCase		)}h{int(_UpperCAmelCase		)}m{int(_UpperCAmelCase		)}s'''

 @property
 def       lowerCAmelCase_							(			self	:	Any		):
       return {"type": "header", "text": {"type": "plain_text", "text": self.title}}

 @property
 def       lowerCAmelCase_							(			self	:	Dict		):
       return {
           "type": "section",
           "text": {
               "type": "plain_text",
               "text": F'''🌞 There were no failures: all {self.n_tests} tests passed. The suite ran in {self.time}.''',
               "emoji": True,
           },
           "accessory": {
               "type": "button",
               "text": {"type": "plain_text", "text": "Check Action results", "emoji": True},
               "url": F'''https://github.com/huggingface/transformers/actions/runs/{os.environ["GITHUB_RUN_ID"]}''',
           },
       }

 @property
 def       lowerCAmelCase_							(			self	:	List[str]		):
       return {
           "type": "section",
           "text": {
               "type": "plain_text",
               "text": (
                   F'''There were {self.n_failures} failures, out of {self.n_tests} tests.\nThe suite ran in'''
                   F''' {self.time}.'''
               ),
               "emoji": True,
           },
           "accessory": {
               "type": "button",
               "text": {"type": "plain_text", "text": "Check Action results", "emoji": True},
               "url": F'''https://github.com/huggingface/transformers/actions/runs/{os.environ["GITHUB_RUN_ID"]}''',
           },
       }

 @property
 def       lowerCAmelCase_							(			self	:	List[str]		):
       _A        =       40
       _A        =       {k: v['failed'] for k, v in doc_test_results.items() if isinstance(_UpperCAmelCase   ,							_UpperCAmelCase		)}

       _A        =       ''
       for category, failures in category_failures.items():
             if len(_UpperCAmelCase		) == 0:
                   continue

             if report != "":
                   report += "\n\n"

             report += F'''*{category} failures*:'''.ljust(line_length // 2		).rjust(line_length // 2		) + "\n"
             report += "`"
             report += "`\n`".join(_UpperCAmelCase		)
             report += "`"

       return {
           "type": "section",
           "text": {
               "type": "mrkdwn",
               "text": F'''The following examples had failures:\n\n\n{report}\n''',
           },
       }

 @property
 def       lowerCAmelCase_							(			self	:	List[Any]		):
       _A        =       [self.header]

       if self.n_failures > 0:
             blocks.append(self.failures		)

       if self.n_failures > 0:
             blocks.extend([self.category_failures]		)

       if self.n_failures == 0:
             blocks.append(self.no_failures		)

       return json.dumps(_UpperCAmelCase		)

 @staticmethod
 def       lowerCAmelCase_							(			):
       _A        =       [
           {
               'type': 'section',
               'text': {
                   'type': 'plain_text',
                   'text': 'There was an issue running the tests.',
               },
               'accessory': {
                   'type': 'button',
                   'text': {'type': 'plain_text', 'text': 'Check Action results', 'emoji': True},
                   'url': F'''https://github.com/huggingface/transformers/actions/runs/{os.environ["GITHUB_RUN_ID"]}''',
               },
           }
       ]

       print('Sending the following payload'		)
       print(json.dumps({'blocks': json.loads(_UpperCAmelCase		)}		)		)

       client.chat_postMessage(
           channel=os.environ['CI_SLACK_CHANNEL_ID_DAILY']   ,							text='There was an issue running the tests.'   ,							blocks=_UpperCAmelCase   ,							)

 def       lowerCAmelCase_							(			self	:	List[str]		):
       print('Sending the following payload'		)
       print(json.dumps({'blocks': json.loads(self.payload		)}		)		)

       _A        =       F'''{self.n_failures} failures out of {self.n_tests} tests,''' if self.n_failures else 'All tests passed.'

       _A        =       client.chat_postMessage(
           channel=os.environ['CI_SLACK_CHANNEL_ID_DAILY']   ,							blocks=self.payload   ,							text=_UpperCAmelCase   ,							)

 def       lowerCAmelCase_							(			self	:	str   ,							_UpperCAmelCase	:	Optional[int]   ,							_UpperCAmelCase	:	List[str]   ,							_UpperCAmelCase	:	str   ,							_UpperCAmelCase	:	Tuple		):
       _A        =       ''
       for key, value in failures.items():
             _A        =       value[:200] + ' [Truncated]' if len(_UpperCAmelCase		) > 250 else value
             failures_text += F'''*{key}*\n_{value}_\n\n'''

       _A        =       job_name
       _A        =       {'type': 'section', 'text': {'type': 'mrkdwn', 'text': text}}

       if job_link is not None:
             _A        =       {
                 'type': 'button',
                 'text': {'type': 'plain_text', 'text': 'GitHub Action job', 'emoji': True},
                 'url': job_link,
             }

       return [
           {"type": "header", "text": {"type": "plain_text", "text": title.upper(), "emoji": True}},
           content,
           {"type": "section", "text": {"type": "mrkdwn", "text": failures_text}},
       ]



 def       lowerCAmelCase_							(			self	:	List[str]		):
       if self.thread_ts is None:
             raise ValueError('Can only post reply if a post has been made.'		)

       _A        =       self.doc_test_results.pop('job_link'		)
       self.doc_test_results.pop('failures'		)
       self.doc_test_results.pop('success'		)
       self.doc_test_results.pop('time_spent'		)

       _A        =       sorted(self.doc_test_results.items()   ,							key=lambda _UpperCAmelCase		: t[0]		)
       for job, job_result in sorted_dict:
             if len(job_result['failures']		):
                   _A        =       F'''*Num failures* :{len(job_result["failed"]		)} \n'''
                   _A        =       job_result['failures']
                   _A        =       self.get_reply_blocks(_UpperCAmelCase   ,							_UpperCAmelCase   ,							_UpperCAmelCase   ,							text=_UpperCAmelCase		)

                   print('Sending the following reply'		)
                   print(json.dumps({'blocks': blocks}		)		)

                   client.chat_postMessage(
                       channel=os.environ['CI_SLACK_CHANNEL_ID_DAILY']   ,							text=F'''Results for {job}'''   ,							blocks=_UpperCAmelCase   ,							thread_ts=self.thread_ts['ts']   ,							)

                   time.sleep(1		)






def 						_snake_case					( )   -> Union[str, Any]:






      '''simple docstring'''


      _A        =       os.environ['GITHUB_RUN_ID']
      _A        =       F'''https://api.github.com/repos/huggingface/transformers/actions/runs/{run_id}/jobs?per_page=100'''
      _A        =       requests.get(_snake_case							).json()
      _A        =       {}

      try:
            jobs.update({job['name']: job['html_url'] for job in result['jobs']}							)
            _A        =       math.ceil((result['total_count'] - 1_00) / 1_00							)

            for i in range(_snake_case							):
                  _A        =       requests.get(url + F'''&page={i + 2}'''							).json()
                  jobs.update({job['name']: job['html_url'] for job in result['jobs']}							)

            return jobs
      except Exception as e:
            print('Unknown error, could not fetch links.'    ,  _snake_case							)

      return {}






def 						_snake_case					( _snake_case							: str							)   -> int:






      '''simple docstring'''


      _A        =       {}

      if os.path.exists(_snake_case							):
            _A        =       os.listdir(_snake_case							)
            for file in files:
                  try:
                        with open(os.path.join(_snake_case    ,  _snake_case							)    ,  encoding='utf-8'							) as f:
                              _A        =       f.read()
                  except UnicodeDecodeError as e:
                        raise ValueError(F'''Could not open {os.path.join(_snake_case    ,  _snake_case							)}.'''							) from e

      return _artifact






def 						_snake_case					( )   -> int:






      '''simple docstring'''


      class     lowercase_	:



       '''simple docstring'''

       def __init__(			self	:	Optional[int]   ,							_UpperCAmelCase	:	str		):
             _A        =       name
             _A        =       []

       def __str__(			self	:	int		):
             return self.name



       def       lowerCAmelCase_							(			self	:	str   ,							_UpperCAmelCase	:	str		):
             self.paths.append({'name': self.name, 'path': path}		)



      _A        =       {}

      _A        =       filter(os.path.isdir    ,  os.listdir()							)
      for directory in directories:
            _A        =       directory
            if artifact_name not in _available_artifacts:
                  _A        =       Artifact(_snake_case							)

                  _available_artifacts[artifact_name].add_path(_snake_case							)

      return _available_artifacts


if __name__ == "__main__":
       a         =      get_job_links()
       a         =      retrieve_available_artifacts()

       a         =      collections.OrderedDict(
           [
               ('''*.py''', '''API Examples'''),
               ('''*.md''', '''MD Examples'''),
           ]
       )

       # This dict will contain all the information relative to each doc test category:
       # - failed: list of failed tests
       # - failures: dict in the format 'test': 'error_message'
       a         =      {
           v: {
               '''failed''': [],
               '''failures''': {},
           }
           for v in docs.values()
       }

       # Link to the GitHub Action job
       a         =      github_actions_job_links.get('''run_doctests''')

       a         =      available_artifacts['''doc_tests_gpu_test_reports'''].paths[0]
       a         =      retrieve_artifact(artifact_path['''name'''])
       if "stats" in artifact:
              a     ,						a     ,						a         =      handle_test_results(artifact['''stats'''])
              a         =      failed
              a         =      success
              a         =      time_spent[1:-1] + ''', '''

              a         =      extract_first_line_failure(artifact['''failures_short'''])
              for line in artifact["summary_short"].split('''\n'''):
                     if re.search('''FAILED''', line):
                            a         =      line.replace('''FAILED ''', '''''')
                            a         =      line.split()[0].replace('''\n''', '''''')

                            if "::" in line:
                                   a     ,						a         =      line.split('''::''')
                            else:
                                   a     ,						a         =      line, line

                            for file_regex in docs.keys():
                                   if fnmatch(file_path, file_regex):
                                          a         =      docs[file_regex]
                                          doc_test_results[category]["failed"].append(test)

                                          a         =      all_failures[test] if test in all_failures else '''N/A'''
                                          a         =      failure
                                          break

       a         =      Message('''🤗 Results of the doc tests.''', doc_test_results)
       message.post()
       message.post_reply()







"""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

a         =      logging.get_logger(__name__)






def 						_snake_case					( _snake_case							: bool    ,  _snake_case							: bool							)   -> Tuple:






      '''simple docstring'''


      def run_func(_snake_case							: Any							):
            @wraps(_snake_case							)
            def run_in_eager_mode(*_snake_case							: List[str]    ,  **_snake_case							: Tuple							):
                  return func(*_snake_case    ,  **_snake_case							)

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

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






      '''simple docstring'''


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







class     lowercase_	(					__lowerCAmelCase							):



 '''simple docstring'''
 UpperCAmelCase   :					TensorFlowBenchmarkArguments
 UpperCAmelCase   :					PretrainedConfig
 UpperCAmelCase   :					str				  =     "TensorFlow"

 @property
 def       lowerCAmelCase_							(			self	:	str		):
       return tf.__version__

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

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

 def       lowerCAmelCase_							(			self	:	Dict   ,							_UpperCAmelCase	:	str   ,							_UpperCAmelCase	:	int   ,							_UpperCAmelCase	:	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]   ,							_UpperCAmelCase		)
       _A        =       self.args.strategy
       if strategy is None:
             raise ValueError('A device strategy has to be initialized before using TensorFlow.'		)
       _A        =       self._prepare_inference_func(_UpperCAmelCase   ,							_UpperCAmelCase   ,							_UpperCAmelCase		)
       return self._measure_memory(_inference		)

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

       _A        =       self._prepare_train_func(_UpperCAmelCase   ,							_UpperCAmelCase   ,							_UpperCAmelCase		)
       return self._measure_memory(_train		)

 def       lowerCAmelCase_							(			self	:	List[Any]   ,							_UpperCAmelCase	:	str   ,							_UpperCAmelCase	:	int   ,							_UpperCAmelCase	:	int		):
       _A        =       self.config_dict[model_name]

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

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

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

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

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

       _A        =       encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

       return _inference

 def       lowerCAmelCase_							(			self	:	str   ,							_UpperCAmelCase	:	str   ,							_UpperCAmelCase	:	int   ,							_UpperCAmelCase	:	int		):
       _A        =       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.'		)

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

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

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

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

       _A        =       encoder_decoder_train if config.is_encoder_decoder else encoder_train

       return _train

 def       lowerCAmelCase_							(			self	:	Tuple   ,							_UpperCAmelCase	:	int		):
       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(_UpperCAmelCase   ,							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
                   _A        =       timeit.repeat(
                       _UpperCAmelCase   ,							repeat=self.args.repeat   ,							number=10   ,							)

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



 def       lowerCAmelCase_							(			self	:	Optional[int]   ,							_UpperCAmelCase	:	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.'		)
                         _A        =       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.'		)
                               _A        =       '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()
                               _A        =       nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx		)
                               _A        =       nvml.nvmlDeviceGetMemoryInfo(_UpperCAmelCase		)
                               _A        =       meminfo.used
                               _A        =       Memory(_UpperCAmelCase		)
                               # 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.'		)
                               _A        =       None
                         else:
                               _A        =       measure_peak_memory_cpu(_UpperCAmelCase		)
                               _A        =       Memory(_UpperCAmelCase		) if isinstance(_UpperCAmelCase   ,							_UpperCAmelCase		) else memory_bytes
                   if self.args.trace_memory_line_by_line:
                         _A        =       stop_memory_tracing(_UpperCAmelCase		)
                         if memory is None:
                               _A        =       summary.total
                   else:
                         _A        =       None

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

