

"""simple docstring"""






from dataclasses import dataclass, field
from typing import Tuple

from ..utils import cached_property, is_tf_available, logging, requires_backends
from .benchmark_args_utils import BenchmarkArguments


if is_tf_available():
 import tensorflow as tf


_lowerCAmelCase						=		logging.get_logger(__name__)







@dataclass
class 	__UpperCamelCase    (    a__						):
  _UpperCAmelCase          =		[
      "no_inference",
      "no_cuda",
      "no_tpu",
      "no_speed",
      "no_memory",
      "no_env_print",
      "no_multi_process",
  ]

  def __init__(		self    ,**_A      ):
      '''simple docstring'''

      for deprecated_arg in self.deprecated_args:
          if deprecated_arg in kwargs:
              _lowerCAmelCase       :   str  =    deprecated_arg[3:]
              _lowerCAmelCase       :   List[Any]  =    not kwargs.pop(_A      )
              logger.warning(
                  F"""{deprecated_arg} is depreciated. Please use --no-{positive_arg} or"""
                  F""" {positive_arg}={kwargs[positive_arg]}"""      )
      _lowerCAmelCase       :   Tuple  =    kwargs.pop('tpu_name'    ,self.tpu_name      )
      _lowerCAmelCase       :   Any  =    kwargs.pop('device_idx'    ,self.device_idx      )
      _lowerCAmelCase       :   List[Any]  =    kwargs.pop('eager_mode'    ,self.eager_mode      )
      _lowerCAmelCase       :   List[str]  =    kwargs.pop('use_xla'    ,self.use_xla      )
      super().__init__(**_A      )

  _UpperCAmelCase          =		field(
      default=a__							,							metadata={"help": "Name of TPU"}							,							)
  _UpperCAmelCase          =		field(
      default=0							,							metadata={"help": "CPU / GPU device index. Defaults to 0."}							,							)
  _UpperCAmelCase          =		field(default=a__							,							metadata={"help": "Benchmark models in eager model."}						)
  _UpperCAmelCase          =		field(
      default=a__							,							metadata={
          "help": "Benchmark models using XLA JIT compilation. Note that `eager_model` has to be set to `False`."
      }							,							)

  @cached_property
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      requires_backends(self    ,['tf']      )
      _lowerCAmelCase       :   Optional[Any]  =    None
      if self.tpu:
          try:
              if self.tpu_name:
                  _lowerCAmelCase       :   Tuple  =    tf.distribute.cluster_resolver.TPUClusterResolver(self.tpu_name      )
              else:
                  _lowerCAmelCase       :   List[str]  =    tf.distribute.cluster_resolver.TPUClusterResolver()
          except ValueError:
              _lowerCAmelCase       :   Union[str, Any]  =    None
      return tpu

  @cached_property
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      requires_backends(self    ,['tf']      )
      if self.is_tpu:
          tf.config.experimental_connect_to_cluster(self._setup_tpu      )
          tf.tpu.experimental.initialize_tpu_system(self._setup_tpu      )

          _lowerCAmelCase       :   Any  =    tf.distribute.TPUStrategy(self._setup_tpu      )
      else:
          # currently no multi gpu is allowed
          if self.is_gpu:
              # TODO: Currently only single GPU is supported
              tf.config.set_visible_devices(self.gpu_list[self.device_idx]    ,'GPU'      )
              _lowerCAmelCase       :   str  =    tf.distribute.OneDeviceStrategy(device=F"""/gpu:{self.device_idx}"""      )
          else:
              tf.config.set_visible_devices([]    ,'GPU'      )  # disable GPU
              _lowerCAmelCase       :   List[str]  =    tf.distribute.OneDeviceStrategy(device=F"""/cpu:{self.device_idx}"""      )

      return strategy

  @property
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      requires_backends(self    ,['tf']      )
      return self._setup_tpu is not None

  @property
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      requires_backends(self    ,['tf']      )
      return self._setup_strategy

  @property
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      requires_backends(self    ,['tf']      )
      return tf.config.list_physical_devices('GPU'      )

  @property
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      requires_backends(self    ,['tf']      )
      if self.cuda:
          return len(self.gpu_list      )
      return 0







  @property
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      return self.n_gpu > 0








"""simple docstring"""






import inspect
import os
import unittest

import torch

import accelerate
from accelerate import Accelerator
from accelerate.test_utils import execute_subprocess_async, require_multi_gpu
from accelerate.utils import patch_environment







class 	__UpperCamelCase    (    unittest.TestCase						):

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Dict  =    inspect.getfile(accelerate.test_utils      )
      _lowerCAmelCase       :   Union[str, Any]  =    os.path.sep.join(mod_file.split(os.path.sep      )[:-1] + ['scripts', 'test_script.py']      )
      _lowerCAmelCase       :   int  =    os.path.sep.join(
          mod_file.split(os.path.sep      )[:-1] + ['scripts', 'test_distributed_data_loop.py']      )
      _lowerCAmelCase       :   Optional[Any]  =    os.path.sep.join(mod_file.split(os.path.sep      )[:-1] + ['scripts', 'test_ops.py']      )

  @require_multi_gpu
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      print(F"""Found {torch.cuda.device_count()} devices."""      )
      _lowerCAmelCase       :   int  =    ['torchrun', F"""--nproc_per_node={torch.cuda.device_count()}""", self.test_file_path]
      with patch_environment(omp_num_threads=1      ):
          execute_subprocess_async(_A    ,env=os.environ.copy()      )

  @require_multi_gpu
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      print(F"""Found {torch.cuda.device_count()} devices."""      )
      _lowerCAmelCase       :   str  =    ['torchrun', F"""--nproc_per_node={torch.cuda.device_count()}""", self.operation_file_path]
      print(F"""Command: {cmd}"""      )
      with patch_environment(omp_num_threads=1      ):
          execute_subprocess_async(_A    ,env=os.environ.copy()      )

  @require_multi_gpu
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      _lowerCAmelCase       :   List[Any]  =    ['torchrun', F"""--nproc_per_node={torch.cuda.device_count()}""", inspect.getfile(self.__class__      )]
      with patch_environment(omp_num_threads=1      ):
          execute_subprocess_async(_A    ,env=os.environ.copy()      )







  @require_multi_gpu
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      print(F"""Found {torch.cuda.device_count()} devices, using 2 devices only"""      )
      _lowerCAmelCase       :   Tuple  =    ['torchrun', F"""--nproc_per_node={torch.cuda.device_count()}""", self.data_loop_file_path]
      with patch_environment(omp_num_threads=1    ,cuda_visible_devices='0,1'      ):
          execute_subprocess_async(_A    ,env=os.environ.copy()      )


if __name__ == "__main__":
 _lowerCAmelCase						=		Accelerator()
 _lowerCAmelCase						=		(accelerator.state.process_index + 2, 1_0)
 _lowerCAmelCase						=		torch.randint(0, 1_0, shape).to(accelerator.device)

 _lowerCAmelCase						=		""""""

 _lowerCAmelCase						=		accelerator.pad_across_processes(tensor)
 if tensora.shape[0] != accelerator.state.num_processes + 1:
  error_msg += F"Found shape {tensora.shape} but should have {accelerator.state.num_processes + 1} at dim 0."
 if not torch.equal(tensora[: accelerator.state.process_index + 2], tensor):
  error_msg += "Tensors have different values."
 if not torch.all(tensora[accelerator.state.process_index + 2 :] == 0):
  error_msg += "Padding was not done with the right value (0)."

 _lowerCAmelCase						=		accelerator.pad_across_processes(tensor, pad_first=True)
 if tensora.shape[0] != accelerator.state.num_processes + 1:
  error_msg += F"Found shape {tensora.shape} but should have {accelerator.state.num_processes + 1} at dim 0."
 _lowerCAmelCase						=		accelerator.state.num_processes - accelerator.state.process_index - 1
 if not torch.equal(tensora[index:], tensor):
  error_msg += "Tensors have different values."
 if not torch.all(tensora[:index] == 0):
  error_msg += "Padding was not done with the right value (0)."

 # Raise error at the end to make sure we don't stop at the first failure.
 if len(error_msg) > 0:
  raise ValueError(error_msg)





