"""simple docstring"""






import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch



def   a_						(					lowerCamelCase       ):  # picklable for multiprocessing
				return x.sum()



def   a_						(					lowerCamelCase       ):  # picklable for multiprocessing
				return i + 1
@dataclass
class 				snake_case       :




					"""simple docstring"""

					snake_case__				    = 42
					snake_case__				    = 42






class 				snake_case       ( __UpperCAmelCase				):




					"""simple docstring"""

					def 		__lowerCAmelCase						(				self   :	List[Any]			):
									UpperCAmelCase__             =							{}
									UpperCAmelCase__             =							[]
									UpperCAmelCase__             =							1
									UpperCAmelCase__             =							[1, 2]
									UpperCAmelCase__             =							{'a': 1, 'b': 2}
									UpperCAmelCase__             =							{'a': [1, 2], 'b': [3, 4]}
									UpperCAmelCase__             =							{'a': {'1': 1}, 'b': 2}
									UpperCAmelCase__             =							{'a': 1, 'b': 2, 'c': 3, 'd': 4}
									UpperCAmelCase__             =							{}
									UpperCAmelCase__             =							[]
									UpperCAmelCase__             =							2
									UpperCAmelCase__             =							[2, 3]
									UpperCAmelCase__             =							{'a': 2, 'b': 3}
									UpperCAmelCase__             =							{'a': [2, 3], 'b': [4, 5]}
									UpperCAmelCase__             =							{'a': {'1': 2}, 'b': 3}
									UpperCAmelCase__             =							{'a': 2, 'b': 3, 'c': 4, 'd': 5}
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__			)    ,lowerCamelCase__			)

									UpperCAmelCase__             =							2
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__    ,num_proc=lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__    ,num_proc=lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__    ,num_proc=lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__    ,num_proc=lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__    ,num_proc=lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__    ,num_proc=lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__    ,num_proc=lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__    ,num_proc=lowerCamelCase__			)    ,lowerCamelCase__			)

									UpperCAmelCase__             =							{'a': np.eye(2			), 'b': np.zeros(3			), 'c': np.ones(2			)}
									UpperCAmelCase__             =							{'a': 2, 'b': 0, 'c': 2}
									UpperCAmelCase__             =							{
									    'a': np.eye(2			).astype(lowerCamelCase__			),
									    'b': np.zeros(3			).astype(lowerCamelCase__			),
									    'c': np.ones(2			).astype(lowerCamelCase__			),
									}
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__    ,map_numpy=lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(
									    {k: v.tolist() for k, v in map_nested(lowerCamelCase__    ,lowerCamelCase__    ,map_numpy=lowerCamelCase__			).items()}    ,{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}    ,)
									self.assertEqual(map_nested(lowerCamelCase__    ,lowerCamelCase__    ,map_numpy=lowerCamelCase__    ,num_proc=lowerCamelCase__			)    ,lowerCamelCase__			)
									self.assertEqual(
									    {k: v.tolist() for k, v in map_nested(lowerCamelCase__    ,lowerCamelCase__    ,map_numpy=lowerCamelCase__    ,num_proc=lowerCamelCase__			).items()}    ,{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}    ,)
									with self.assertRaises(lowerCamelCase__			):  # can't pickle a local lambda
													map_nested(lambda lowerCamelCase__			: x + 1    ,lowerCamelCase__    ,num_proc=lowerCamelCase__			)







					def 		__lowerCAmelCase						(				self   :	Any			):
									UpperCAmelCase__             =							{'a': 1, 'b': 2}
									UpperCAmelCase__             =							{'a': 3, 'b': 4}
									UpperCAmelCase__             =							{'a': 5, 'b': 6}
									UpperCAmelCase__             =							sorted([('a', (1, 3, 5)), ('b', (2, 4, 6))]			)
									self.assertEqual(sorted(zip_dict(lowerCamelCase__    ,lowerCamelCase__    ,lowerCamelCase__			)			)    ,lowerCamelCase__			)







					def 		__lowerCAmelCase						(				self   :	Optional[Any]			):
									class 				snake_case       :




														"""simple docstring"""

														snake_case__				    = "bar"






									UpperCAmelCase__             =							Foo()
									self.assertEqual(foo.my_attr    ,'bar'			)
									with temporary_assignment(lowerCamelCase__    ,'my_attr'    ,'BAR'			):
													self.assertEqual(foo.my_attr    ,'BAR'			)
									self.assertEqual(foo.my_attr    ,'bar'			)






@pytest.mark.parametrize(
    'iterable_length, num_proc, expected_num_proc'					,			[
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (1_6, 1_6, 1_6),
        (1_6, 1_7, 1_6),
        (1_7, 1_6, 1_6),
    ]					,			)
def   a_						(					lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase       ):
				with patch('datasets.utils.py_utils._single_map_nested'       ) as mock_single_map_nested, patch(
				    'datasets.parallel.parallel.Pool'       ) as mock_multiprocessing_pool:
								UpperCAmelCase__             =							{f'''{i}''': i for i in range(lowerCamelCase       )}
								UpperCAmelCase__             =							map_nested(lambda lowerCamelCase       : x + 1_0					,			lowerCamelCase					,			num_proc=lowerCamelCase					,			parallel_min_length=1_6       )
								if expected_num_proc == 1:
												assert mock_single_map_nested.called
												assert not mock_multiprocessing_pool.called
								else:
												assert not mock_single_map_nested.called
												assert mock_multiprocessing_pool.called
												assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc
class 				snake_case       ( __UpperCAmelCase				):




					"""simple docstring"""

					@require_tf
					def 		__lowerCAmelCase						(				self   :	Any			):
									import tensorflow as tf
									from tensorflow.keras import layers

									UpperCAmelCase__             =							layers.Dense(2			)

									def gen_random_output():
													UpperCAmelCase__             =							tf.random.uniform((1, 3)			)
													return model(lowerCamelCase__			).numpy()

									with temp_seed(42    ,set_tensorflow=lowerCamelCase__			):
													UpperCAmelCase__             =							gen_random_output()
									with temp_seed(42    ,set_tensorflow=lowerCamelCase__			):
													UpperCAmelCase__             =							gen_random_output()
									UpperCAmelCase__             =							gen_random_output()

									np.testing.assert_equal(lowerCamelCase__    ,lowerCamelCase__			)
									self.assertGreater(np.abs(outa - outa			).sum()    ,0			)







					@require_torch
					def 		__lowerCAmelCase						(				self   :	Dict			):
									import torch

									def gen_random_output():
													UpperCAmelCase__             =							torch.nn.Linear(3    ,2			)
													UpperCAmelCase__             =							torch.rand(1    ,3			)
													return model(lowerCamelCase__			).detach().numpy()

									with temp_seed(42    ,set_pytorch=lowerCamelCase__			):
													UpperCAmelCase__             =							gen_random_output()
									with temp_seed(42    ,set_pytorch=lowerCamelCase__			):
													UpperCAmelCase__             =							gen_random_output()
									UpperCAmelCase__             =							gen_random_output()

									np.testing.assert_equal(lowerCamelCase__    ,lowerCamelCase__			)
									self.assertGreater(np.abs(outa - outa			).sum()    ,0			)







					def 		__lowerCAmelCase						(				self   :	str			):
									def gen_random_output():
													return np.random.rand(1    ,3			)

									with temp_seed(42			):
													UpperCAmelCase__             =							gen_random_output()
									with temp_seed(42			):
													UpperCAmelCase__             =							gen_random_output()
									UpperCAmelCase__             =							gen_random_output()

									np.testing.assert_equal(lowerCamelCase__    ,lowerCamelCase__			)
									self.assertGreater(np.abs(outa - outa			).sum()    ,0			)






@pytest.mark.parametrize('input_data'					,			[{}]       )
def   a_						(					lowerCamelCase       ):
				UpperCAmelCase__             =							NestedDataStructure(lowerCamelCase       ).data
				assert output_data == input_data



@pytest.mark.parametrize(
    'data, expected_output'					,			[
        ({}, []),
        ([], []),
        ('foo', ['foo']),
        (['foo', 'bar'], ['foo', 'bar']),
        ([['foo', 'bar']], ['foo', 'bar']),
        ([[['foo'], ['bar']]], ['foo', 'bar']),
        ([[['foo'], 'bar']], ['foo', 'bar']),
        ({'a': 1, 'b': 2}, [1, 2]),
        ({'a': [1, 2], 'b': [3, 4]}, [1, 2, 3, 4]),
        ({'a': [[1, 2]], 'b': [[3, 4]]}, [1, 2, 3, 4]),
        ({'a': [[1, 2]], 'b': [3, 4]}, [1, 2, 3, 4]),
        ({'a': [[[1], [2]]], 'b': [[[3], [4]]]}, [1, 2, 3, 4]),
        ({'a': [[[1], [2]]], 'b': [[3, 4]]}, [1, 2, 3, 4]),
        ({'a': [[[1], [2]]], 'b': [3, 4]}, [1, 2, 3, 4]),
        ({'a': [[[1], [2]]], 'b': [3, [4]]}, [1, 2, 3, 4]),
        ({'a': {'1': 1}, 'b': 2}, [1, 2]),
        ({'a': {'1': [1]}, 'b': 2}, [1, 2]),
        ({'a': {'1': [1]}, 'b': [2]}, [1, 2]),
    ]					,			)
def   a_						(					lowerCamelCase					,			lowerCamelCase       ):
				UpperCAmelCase__             =							NestedDataStructure(lowerCamelCase       ).flatten()
				assert output == expected_output



def   a_						(					):
				UpperCAmelCase__             =							A(x=1					,			y='foobar'       )
				UpperCAmelCase__             =							{'x': 1, 'y': 'foobar'}
				assert asdict(lowerCamelCase       ) == expected_output

				UpperCAmelCase__             =							{'a': {'b': A(x=1_0					,			y='foo'       )}, 'c': [A(x=2_0					,			y='bar'       )]}
				UpperCAmelCase__             =							{'a': {'b': {'x': 1_0, 'y': 'foo'}}, 'c': [{'x': 2_0, 'y': 'bar'}]}
				assert asdict(lowerCamelCase       ) == expected_output

				with pytest.raises(lowerCamelCase       ):
								asdict([1, A(x=1_0					,			y='foo'       )]       )



def   a_						(					lowerCamelCase       ):
				return text.split()



def   a_						(					lowerCamelCase       ):
				yield (time.time(), content)
				time.sleep(2       )
				yield (time.time(), content)







def   a_						(					):
				with Pool(2       ) as pool:
								UpperCAmelCase__             =							list(iflatmap_unordered(lowerCamelCase					,			_split_text					,			kwargs_iterable=[{'text': 'hello there'}] * 1_0       )       )
								assert out.count('hello'       ) == 1_0
								assert out.count('there'       ) == 1_0
								assert len(lowerCamelCase       ) == 2_0

				# check multiprocess from pathos (uses dill for pickling)
				with multiprocess.Pool(2       ) as pool:
								UpperCAmelCase__             =							list(iflatmap_unordered(lowerCamelCase					,			_split_text					,			kwargs_iterable=[{'text': 'hello there'}] * 1_0       )       )
								assert out.count('hello'       ) == 1_0
								assert out.count('there'       ) == 1_0
								assert len(lowerCamelCase       ) == 2_0

				# check that we get items as fast as possible
				with Pool(2       ) as pool:
								UpperCAmelCase__             =							[]
								for yield_time, content in iflatmap_unordered(
								    lowerCamelCase					,			_aseconds_generator_of_aitems_with_timing					,			kwargs_iterable=[{'content': 'a'}, {'content': 'b'}]       ):
												assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
												out.append(lowerCamelCase       )
								assert out.count('a'       ) == 2
								assert out.count('b'       ) == 2
								assert len(lowerCamelCase       ) == 4




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

UpperCAmelCase__     						=  logging.get_logger(__name__)






def 			__UpperCAmelCase    (	lowercase      ,lowercase						):







  """simple docstring"""
  def run_func(lowercase						):
    @wraps(lowercase						)
    def run_in_eager_mode(*lowercase      ,**lowercase						):
      return func(*lowercase      ,**lowercase						)

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

    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    (	lowercase      ,lowercase      ,lowercase						):







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


class 				a (						lowerCAmelCase_      ):
    _snake_case   :     TensorFlowBenchmarkArguments
    _snake_case   :     PretrainedConfig
    _snake_case   :     str   							=							"TensorFlow"


    @property
    def    lowerCAmelCase_     (    self		:     Union[str, Any]    ):
      return tf.__version__


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


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


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


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

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


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

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

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

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

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

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

      _UpperCAmelCase         =					encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

      return _inference


    def    lowerCAmelCase_     (    self		:     List[str]      , __lowerCAmelCase		:     str      , __lowerCAmelCase		:     int      , __lowerCAmelCase		:     int    ):
      _UpperCAmelCase         =					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."""    )

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

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

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

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

      _UpperCAmelCase         =					encoder_decoder_train if config.is_encoder_decoder else encoder_train

      return _train


    def    lowerCAmelCase_     (    self		:     Union[str, Any]      , __lowerCAmelCase		:     Any    ):
      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(__lowerCAmelCase      , 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
          _UpperCAmelCase         =					timeit.repeat(
              __lowerCAmelCase      , repeat=self.args.repeat      , number=10      , )

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


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

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