





import gc
import math
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, logging, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


_A :   Optional[Any]  			=  logging.get_logger(__name__)

enable_full_determinism()


class        __SCREAMING_SNAKE_CASE      (			lowerCAmelCase_					,lowerCAmelCase_					,unittest.TestCase ):






       """simple docstring"""




       _UpperCAmelCase				:  int			   =   UNetaDModel
       _UpperCAmelCase				:  List[Any]			   =   "sample"







       @property
       def 						__lowerCamelCase   (				self					:     List[str]      )       ->Any:
          lowerCamelCase__       :						List[Any]        =							4
          lowerCamelCase__       :						List[str]        =							3
          lowerCamelCase__       :						Dict        =							(3_2, 3_2)

          lowerCamelCase__       :						Optional[Any]        =							floats_tensor((batch_size, num_channels) + sizes      ).to(A      )
          lowerCamelCase__       :						List[Any]        =							torch.tensor([1_0]      ).to(A      )

          return {"sample": noise, "timestep": time_step}







       @property
       def 						__lowerCamelCase   (				self					:     Optional[int]      )       ->Optional[int]:
          return (3, 3_2, 3_2)







       @property
       def 						__lowerCamelCase   (				self					:     int      )       ->Dict:
          return (3, 3_2, 3_2)







       def 						__lowerCamelCase   (				self					:     Optional[Any]      )       ->Optional[Any]:
          lowerCamelCase__       :						Optional[int]        =							{
              '''block_out_channels''': (3_2, 6_4),
              '''down_block_types''': ('''DownBlock2D''', '''AttnDownBlock2D'''),
              '''up_block_types''': ('''AttnUpBlock2D''', '''UpBlock2D'''),
              '''attention_head_dim''': 3,
              '''out_channels''': 3,
              '''in_channels''': 3,
              '''layers_per_block''': 2,
              '''sample_size''': 3_2,
          }
          lowerCamelCase__       :						Dict        =							self.dummy_input
          return init_dict, inputs_dict


class        __SCREAMING_SNAKE_CASE      (			lowerCAmelCase_					,lowerCAmelCase_					,unittest.TestCase ):






       """simple docstring"""




       _UpperCAmelCase				:  Optional[Any]			   =   UNetaDModel
       _UpperCAmelCase				:  Dict			   =   "sample"







       @property
       def 						__lowerCamelCase   (				self					:     List[str]      )       ->int:
          lowerCamelCase__       :						Tuple        =							4
          lowerCamelCase__       :						Dict        =							4
          lowerCamelCase__       :						Dict        =							(3_2, 3_2)

          lowerCamelCase__       :						List[Any]        =							floats_tensor((batch_size, num_channels) + sizes      ).to(A      )
          lowerCamelCase__       :						Optional[int]        =							torch.tensor([1_0]      ).to(A      )

          return {"sample": noise, "timestep": time_step}







       @property
       def 						__lowerCamelCase   (				self					:     Optional[int]      )       ->int:
          return (4, 3_2, 3_2)







       @property
       def 						__lowerCamelCase   (				self					:     Dict      )       ->Union[str, Any]:
          return (4, 3_2, 3_2)







       def 						__lowerCamelCase   (				self					:     List[Any]      )       ->Any:
          lowerCamelCase__       :						Tuple        =							{
              '''sample_size''': 3_2,
              '''in_channels''': 4,
              '''out_channels''': 4,
              '''layers_per_block''': 2,
              '''block_out_channels''': (3_2, 6_4),
              '''attention_head_dim''': 3_2,
              '''down_block_types''': ('''DownBlock2D''', '''DownBlock2D'''),
              '''up_block_types''': ('''UpBlock2D''', '''UpBlock2D'''),
          }
          lowerCamelCase__       :						Tuple        =							self.dummy_input
          return init_dict, inputs_dict







       def 						__lowerCamelCase   (				self					:     List[str]      )       ->Any:
          lowerCamelCase__       :						int        =							UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''		, output_loading_info=A      )

          self.assertIsNotNone(A      )
          self.assertEqual(len(loading_info['''missing_keys''']      )		, 0      )

          model.to(A      )
          lowerCamelCase__       :						Optional[int]        =							model(**self.dummy_input      ).sample

          assert image is not None, "Make sure output is not None"







       @unittest.skipIf(torch_device != '''cuda'''		, '''This test is supposed to run on GPU'''      )
       def 						__lowerCamelCase   (				self					:     str      )       ->str:
          lowerCamelCase__       :						Union[str, Any]        =							UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''		, output_loading_info=A      )
          model.to(A      )
          lowerCamelCase__       :						Optional[Any]        =							model(**self.dummy_input      ).sample

          assert image is not None, "Make sure output is not None"







       @unittest.skipIf(torch_device != '''cuda'''		, '''This test is supposed to run on GPU'''      )
       def 						__lowerCamelCase   (				self					:     Optional[int]      )       ->Dict:
          # by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
          lowerCamelCase__       :						Optional[int]        =							UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''		, output_loading_info=A      )
          model_accelerate.to(A      )
          model_accelerate.eval()

          lowerCamelCase__       :						Tuple        =							torch.randn(
              1		, model_accelerate.config.in_channels		, model_accelerate.config.sample_size		, model_accelerate.config.sample_size		, generator=torch.manual_seed(0      )		, )
          lowerCamelCase__       :						Any        =							noise.to(A      )
          lowerCamelCase__       :						List[Any]        =							torch.tensor([1_0] * noise.shape[0]      ).to(A      )

          lowerCamelCase__       :						Optional[Any]        =							model_accelerate(A		, A      )['''sample''']

          # two models don't need to stay in the device at the same time
          del model_accelerate
          torch.cuda.empty_cache()
          gc.collect()

          lowerCamelCase__       :						str        =							UNetaDModel.from_pretrained(
              '''fusing/unet-ldm-dummy-update'''		, output_loading_info=A		, low_cpu_mem_usage=A      )
          model_normal_load.to(A      )
          model_normal_load.eval()
          lowerCamelCase__       :						Union[str, Any]        =							model_normal_load(A		, A      )['''sample''']

          assert torch_all_close(A		, A		, rtol=1e-3      )







       def 						__lowerCamelCase   (				self					:     str      )       ->Dict:
          lowerCamelCase__       :						int        =							UNetaDModel.from_pretrained('''fusing/unet-ldm-dummy-update'''      )
          model.eval()
          model.to(A      )

          lowerCamelCase__       :						Union[str, Any]        =							torch.randn(
              1		, model.config.in_channels		, model.config.sample_size		, model.config.sample_size		, generator=torch.manual_seed(0      )		, )
          lowerCamelCase__       :						List[Any]        =							noise.to(A      )
          lowerCamelCase__       :						Optional[Any]        =							torch.tensor([1_0] * noise.shape[0]      ).to(A      )

          with torch.no_grad():
             lowerCamelCase__       :						List[str]        =							model(A		, A      ).sample

          lowerCamelCase__       :						str        =							output[0, -1, -3:, -3:].flatten().cpu()
          # fmt: off
          lowerCamelCase__       :						List[str]        =							torch.tensor([-13.32_58, -20.11_00, -15.98_73, -17.66_17, -23.05_96, -17.94_19, -13.36_75, -16.18_89, -12.38_00]      )
          # fmt: on

          self.assertTrue(torch_all_close(A		, A		, rtol=1e-3      )      )


class        __SCREAMING_SNAKE_CASE      (			lowerCAmelCase_					,lowerCAmelCase_					,unittest.TestCase ):






       """simple docstring"""




       _UpperCAmelCase				:  str			   =   UNetaDModel
       _UpperCAmelCase				:  List[str]			   =   "sample"







       @property
       def 						__lowerCamelCase   (				self					:     Tuple		, A					:     Union[str, Any]=(3_2, 3_2)      )       ->Any:
          lowerCamelCase__       :						List[str]        =							4
          lowerCamelCase__       :						Dict        =							3

          lowerCamelCase__       :						Optional[Any]        =							floats_tensor((batch_size, num_channels) + sizes      ).to(A      )
          lowerCamelCase__       :						List[Any]        =							torch.tensor(batch_size * [1_0]      ).to(dtype=torch.intaa		, device=A      )

          return {"sample": noise, "timestep": time_step}







       @property
       def 						__lowerCamelCase   (				self					:     int      )       ->Any:
          return (3, 3_2, 3_2)







       @property
       def 						__lowerCamelCase   (				self					:     Dict      )       ->Tuple:
          return (3, 3_2, 3_2)







       def 						__lowerCamelCase   (				self					:     str      )       ->Optional[Any]:
          lowerCamelCase__       :						Union[str, Any]        =							{
              '''block_out_channels''': [3_2, 6_4, 6_4, 6_4],
              '''in_channels''': 3,
              '''layers_per_block''': 1,
              '''out_channels''': 3,
              '''time_embedding_type''': '''fourier''',
              '''norm_eps''': 1e-6,
              '''mid_block_scale_factor''': math.sqrt(2.0      ),
              '''norm_num_groups''': None,
              '''down_block_types''': [
                  '''SkipDownBlock2D''',
                  '''AttnSkipDownBlock2D''',
                  '''SkipDownBlock2D''',
                  '''SkipDownBlock2D''',
              ],
              '''up_block_types''': [
                  '''SkipUpBlock2D''',
                  '''SkipUpBlock2D''',
                  '''AttnSkipUpBlock2D''',
                  '''SkipUpBlock2D''',
              ],
          }
          lowerCamelCase__       :						Any        =							self.dummy_input
          return init_dict, inputs_dict







       @slow
       def 						__lowerCamelCase   (				self					:     Optional[Any]      )       ->List[Any]:
          lowerCamelCase__       :						str        =							UNetaDModel.from_pretrained('''google/ncsnpp-celebahq-256'''		, output_loading_info=A      )
          self.assertIsNotNone(A      )
          self.assertEqual(len(loading_info['''missing_keys''']      )		, 0      )

          model.to(A      )
          lowerCamelCase__       :						Optional[Any]        =							self.dummy_input
          lowerCamelCase__       :						Union[str, Any]        =							floats_tensor((4, 3) + (2_5_6, 2_5_6)      ).to(A      )
          lowerCamelCase__       :						str        =							noise
          lowerCamelCase__       :						int        =							model(**A      )

          assert image is not None, "Make sure output is not None"







       @slow
       def 						__lowerCamelCase   (				self					:     Tuple      )       ->List[str]:
          lowerCamelCase__       :						Optional[int]        =							UNetaDModel.from_pretrained('''google/ncsnpp-celebahq-256'''      )
          model.to(A      )

          lowerCamelCase__       :						Dict        =							4
          lowerCamelCase__       :						Tuple        =							3
          lowerCamelCase__       :						Dict        =							(2_5_6, 2_5_6)

          lowerCamelCase__       :						Optional[Any]        =							torch.ones((batch_size, num_channels) + sizes      ).to(A      )
          lowerCamelCase__       :						Optional[Any]        =							torch.tensor(batch_size * [1e-4]      ).to(A      )

          with torch.no_grad():
             lowerCamelCase__       :						Tuple        =							model(A		, A      ).sample

          lowerCamelCase__       :						str        =							output[0, -3:, -3:, -1].flatten().cpu()
          # fmt: off
          lowerCamelCase__       :						Optional[int]        =							torch.tensor([-4_8_4_2.8_6_9_1, -6_4_9_9.6_6_3_1, -3_8_0_0.1_9_5_3, -7_9_7_8.2_6_8_6, -1_0_9_8_0.7_1_2_9, -2_0_0_2_8.8_5_3_5, 8_1_4_8.2_8_2_2, 2_3_4_2.2_9_0_5, 5_6_7.7_6_0_8]      )
          # fmt: on

          self.assertTrue(torch_all_close(A		, A		, rtol=1e-2      )      )







       def 						__lowerCamelCase   (				self					:     List[str]      )       ->Optional[int]:
          lowerCamelCase__       :						Tuple        =							UNetaDModel.from_pretrained('''fusing/ncsnpp-ffhq-ve-dummy-update'''      )
          model.to(A      )

          lowerCamelCase__       :						List[Any]        =							4
          lowerCamelCase__       :						int        =							3
          lowerCamelCase__       :						Union[str, Any]        =							(3_2, 3_2)

          lowerCamelCase__       :						Union[str, Any]        =							torch.ones((batch_size, num_channels) + sizes      ).to(A      )
          lowerCamelCase__       :						Dict        =							torch.tensor(batch_size * [1e-4]      ).to(A      )

          with torch.no_grad():
             lowerCamelCase__       :						Union[str, Any]        =							model(A		, A      ).sample

          lowerCamelCase__       :						int        =							output[0, -3:, -3:, -1].flatten().cpu()
          # fmt: off
          lowerCamelCase__       :						str        =							torch.tensor([-0.03_25, -0.09_00, -0.08_69, -0.03_32, -0.07_25, -0.02_70, -0.01_01, 0.02_27, 0.02_56]      )
          # fmt: on

          self.assertTrue(torch_all_close(A		, A		, rtol=1e-2      )      )







       def 						__lowerCamelCase   (				self					:     Optional[Any]      )       ->Any:
          # not required for this model
          pass










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 :   Any  			=  logging.get_logger(__name__)



def 	_a						(   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 	_a						(   UpperCAmelCase      ,  UpperCAmelCase      ,  UpperCAmelCase					)    ->    ["tf.Tensor"]:







   """simple docstring"""
   lowerCamelCase__       :						List[Any]        =							random.Random()
   lowerCamelCase__       :						str        =							[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        __SCREAMING_SNAKE_CASE      (			lowerCAmelCase_ ):
       _UpperCAmelCase				:  TensorFlowBenchmarkArguments
       _UpperCAmelCase				:  PretrainedConfig
       _UpperCAmelCase				:  str			   =   "TensorFlow"







       @property
       def 						__lowerCamelCase   (				self					:     int      )       ->Optional[int]:
          return tf.__version__







       def 						__lowerCamelCase   (				self					:     Optional[int]		, A					:     str		, A					:     int		, A					:     int      )       ->float:
          # initialize GPU on separate process
          lowerCamelCase__       :						Dict        =							self.args.strategy
          if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''      )
          lowerCamelCase__       :						int        =							self._prepare_inference_func(A		, A		, A      )
          return self._measure_speed(_inference      )







       def 						__lowerCamelCase   (				self					:     str		, A					:     str		, A					:     int		, A					:     int      )       ->float:
          lowerCamelCase__       :						Optional[int]        =							self.args.strategy
          if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''      )
          lowerCamelCase__       :						List[Any]        =							self._prepare_train_func(A		, A		, A      )
          return self._measure_speed(_train      )







       def 						__lowerCamelCase   (				self					:     int		, A					:     str		, A					:     int		, A					:     int      )       ->[Memory, Optional[MemorySummary]]:
          # 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__       :						int        =							self.args.strategy
          if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''      )
          lowerCamelCase__       :						str        =							self._prepare_inference_func(A		, A		, A      )
          return self._measure_memory(_inference      )







       def 						__lowerCamelCase   (				self					:     List[str]		, A					:     str		, A					:     int		, A					:     int      )       ->[Memory, Optional[MemorySummary]]:
          if self.args.is_gpu:
             tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]		, A      )
          lowerCamelCase__       :						List[Any]        =							self.args.strategy
          if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''      )

          lowerCamelCase__       :						str        =							self._prepare_train_func(A		, A		, A      )
          return self._measure_memory(_train      )







       def 						__lowerCamelCase   (				self					:     Dict		, A					:     str		, A					:     int		, A					:     int      )       ->Callable[[], None]:
          lowerCamelCase__       :						Tuple        =							self.config_dict[model_name]

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

          lowerCamelCase__       :						Tuple        =							(
              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__       :						Any        =							'''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
                lowerCamelCase__       :						List[Any]        =							__import__('''transformers'''		, fromlist=[model_class]      )
                lowerCamelCase__       :						int        =							getattr(A		, A      )
                lowerCamelCase__       :						int        =							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__       :						Union[str, Any]        =							TF_MODEL_MAPPING[config.__class__](A      )

          # encoder-decoder has vocab size saved differently
          lowerCamelCase__       :						Tuple        =							config.vocab_size if hasattr(A		, '''vocab_size'''      ) else config.encoder.vocab_size
          lowerCamelCase__       :						Optional[Any]        =							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__       :						int        =							encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

          return _inference







       def 						__lowerCamelCase   (				self					:     List[str]		, A					:     str		, A					:     int		, A					:     int      )       ->Callable[[], None]:
          lowerCamelCase__       :						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.'''      )

          lowerCamelCase__       :						Optional[int]        =							(
              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__       :						Any        =							'''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
                lowerCamelCase__       :						List[str]        =							__import__('''transformers'''		, fromlist=[model_class]      )
                lowerCamelCase__       :						Optional[int]        =							getattr(A		, A      )
                lowerCamelCase__       :						Optional[Any]        =							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__       :						List[str]        =							TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](A      )

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

          @run_with_tf_optimizations(self.args.eager_mode		, self.args.use_xla      )
          def encoder_decoder_train():
             lowerCamelCase__       :						int        =							model(A		, decoder_input_ids=A		, labels=A		, training=A      )[0]
             lowerCamelCase__       :						List[Any]        =							tf.gradients(A		, model.trainable_variables      )
             return gradients

          @run_with_tf_optimizations(self.args.eager_mode		, self.args.use_xla      )
          def encoder_train():
             lowerCamelCase__       :						Optional[int]        =							model(A		, labels=A		, training=A      )[0]
             lowerCamelCase__       :						List[str]        =							tf.gradients(A		, model.trainable_variables      )
             return gradients

          lowerCamelCase__       :						Tuple        =							encoder_decoder_train if config.is_encoder_decoder else encoder_train

          return _train







       def 						__lowerCamelCase   (				self					:     Tuple		, A					:     Any      )       ->float:
          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__       :						Optional[Any]        =							timeit.repeat(
                    A		, repeat=self.args.repeat		, number=1_0		, )

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







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

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



