




'''simple docstring'''







import inspect
import warnings

import torch

from .state import AcceleratorState, GradientState
from .utils import DistributedType, honor_type, is_tpu_available


if is_tpu_available(check_device=False):
    import torch_xla.core.xla_model as xm







def 	_A			(   UpperCAmelCase			,UpperCAmelCase	):






   '''simple docstring'''







   if isinstance(UpperCAmelCase			,(list, tuple)	):
      return honor_type(UpperCAmelCase			,(move_to_device(UpperCAmelCase			,UpperCAmelCase	) for t in state)	)
   elif isinstance(UpperCAmelCase			,UpperCAmelCase	):
      return type(UpperCAmelCase	)({k: move_to_device(UpperCAmelCase			,UpperCAmelCase	) for k, v in state.items()}	)
   elif isinstance(UpperCAmelCase			,torch.Tensor	):
      return state.to(UpperCAmelCase	)
   return state







class   _snake_case( torch.optim.Optimizer			):






 def __init__(self					:     List[Any]						,		a					:     Dict						,		a					:     Optional[int]=True						,		a					:     List[str]=None						)					->      int:






    """simple docstring"""


    A__         =    optimizer
    A__         =    scaler
    A__         =    AcceleratorState()
    A__         =    GradientState()
    A__         =    device_placement
    A__         =    False
    A__         =    None

    # Handle device placement
    if device_placement:
       A__         =    self.optimizer.state_dict()
       if self.accelerator_state.distributed_type == DistributedType.TPU:
          xm.send_cpu_data_to_device(a						,		self.accelerator_state.device						)
       else:
          A__         =    move_to_device(a						,		self.accelerator_state.device						)
       self.optimizer.load_state_dict(a						)






 @property
 def 	_UpperCamelCase   (self					:     List[str]						)					->      Optional[int]:






    """simple docstring"""


    return self.optimizer.state






 @state.setter
 def 	_UpperCamelCase   (self					:     List[Any]						,		a					:     Dict						)					->      Tuple:






    """simple docstring"""


    A__         =    state






 @property
 def 	_UpperCamelCase   (self					:     Optional[int]						)					->      Union[str, Any]:






    """simple docstring"""


    return self.optimizer.param_groups






 @param_groups.setter
 def 	_UpperCamelCase   (self					:     Tuple						,		a					:     Any						)					->      List[Any]:






    """simple docstring"""


    A__         =    param_groups






 @property
 def 	_UpperCamelCase   (self					:     List[Any]						)					->      Dict:






    """simple docstring"""


    return self.optimizer.defaults






 @defaults.setter
 def 	_UpperCamelCase   (self					:     Optional[Any]						,		a					:     Dict						)					->      Optional[int]:






    """simple docstring"""


    A__         =    defaults






 def 	_UpperCamelCase   (self					:     Union[str, Any]						,		a					:     Optional[int]						)					->      Optional[int]:






    """simple docstring"""


    self.optimizer.add_param_group(a						)






 def 	_UpperCamelCase   (self					:     Optional[Any]						,		a					:     Any						)					->      str:






    """simple docstring"""


    if self.accelerator_state.distributed_type == DistributedType.TPU and self.device_placement:
       xm.send_cpu_data_to_device(a						,		self.accelerator_state.device						)
    self.optimizer.load_state_dict(a						)






 def 	_UpperCamelCase   (self					:     str						)					->      Optional[int]:






    """simple docstring"""


    return self.optimizer.state_dict()






 def 	_UpperCamelCase   (self					:     Dict						,		a					:     Any=None						)					->      List[str]:






    """simple docstring"""


    if self.gradient_state.sync_gradients:
       A__         =    "set_to_none" in inspect.signature(self.optimizer.zero_grad						).parameters
       if accept_arg:
          if set_to_none is None:
             A__         =    False
          self.optimizer.zero_grad(set_to_none=a						)
       else:
          if set_to_none is not None:
             raise ValueError("`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer."						)
          self.optimizer.zero_grad()






 def 	_UpperCamelCase   (self					:     Union[str, Any]						,		a					:     Dict=None						)					->      Optional[int]:






    """simple docstring"""


    if self.gradient_state.sync_gradients:
       if self.accelerator_state.distributed_type == DistributedType.TPU:
          A__         =    {"closure": closure} if closure is not None else {}
          xm.optimizer_step(self.optimizer						,		optimizer_args=a						)
       elif self.scaler is not None:
          A__         =    False
          if self._last_scale is None:
             # `get_scale` is an async operation requiring full synchronization
             # on CPU and GPUs before finishing. As a result, we store away
             # the prior one to reduce the call overhead
             A__         =    self.scaler.get_scale()
             A__         =    True
          self.scaler.step(self.optimizer						,		a						)
          self.scaler.update()
          A__         =    self.scaler.get_scale()
          if not new_scale:
             # If we reduced the loss scale, it means the optimizer step was skipped because of gradient overflow.
             A__         =    scale_after < self._last_scale
          A__         =    scale_after

       else:
          self.optimizer.step(a						)






 def 	_UpperCamelCase   (self					:     List[str]						,		a					:     Any						)					->      List[Any]:






    """simple docstring"""


    for param_group in self.optimizer.param_groups:
       A__         =    [parameters_map.get(a						,		a						) for p in param_group["params"]]






 @property
 def 	_UpperCamelCase   (self					:     List[str]						)					->      str:






    """simple docstring"""


    warnings.warn(
        "The `is_overflow` property is deprecated and will be removed in version 1.0 of Accelerate use "
        "`optimizer.step_was_skipped` instead."						,		a						,		)
    return self._is_overflow






 @property
 def 	_UpperCamelCase   (self					:     Any						)					->      Any:






    """simple docstring"""


    return self._is_overflow






 def __getstate__(self					:     Any						)					->      str:






    """simple docstring"""


    return self.__dict__.copy()






 def __setstate__(self					:     Tuple						,		a					:     Union[str, Any]						)					->      Tuple:






    """simple docstring"""


    self.__dict__.update(a						)








'''simple docstring'''







import warnings
from dataclasses import dataclass, field
from typing import Optional, Tuple

from .training_args import TrainingArguments
from .utils import cached_property, is_tf_available, logging, requires_backends


lowerCAmelCase_    	=							logging.get_logger(__name__)

if is_tf_available():
    import tensorflow as tf

@dataclass
class   _snake_case( UpperCAmelCase			):
 __snake_case:       Optional[Any]									=				'''tf'''
 __snake_case:       Optional[str]									=				field(
     default=UpperCAmelCase					,   metadata={'''help''': '''Name of TPU'''}					,   )

 __snake_case:       Optional[str]									=				field(
     default=UpperCAmelCase					,   metadata={'''help''': '''Zone of TPU'''}					,   )

 __snake_case:       Optional[str]									=				field(
     default=UpperCAmelCase					,   metadata={'''help''': '''Name of Cloud TPU-enabled project'''}					,   )

 __snake_case:       float									=				field(
     default=1.0					,   metadata={'''help''': '''Power for the Polynomial decay LR scheduler.'''}					,   )

 __snake_case:       bool									=				field(default=UpperCAmelCase					,   metadata={'''help''': '''Whether to activate the XLA compilation or not'''}			)






 @cached_property
 def 	_UpperCamelCase   (self					:     Tuple						)					->      Tuple["tf.distribute.Strategy", int]:






    """simple docstring"""


    requires_backends(self						,		['tf']						)
    logger.info('Tensorflow: setting up strategy'						)

    A__         =    tf.config.list_physical_devices('GPU'						)

    # Set to float16 at first
    if self.fpaa:
       tf.keras.mixed_precision.set_global_policy('mixed_float16'						)

    if self.no_cuda:
       A__         =    tf.distribute.OneDeviceStrategy(device='/cpu:0'						)
    else:
       try:
          if self.tpu_name:
             A__         =    tf.distribute.cluster_resolver.TPUClusterResolver(
                 self.tpu_name						,		zone=self.tpu_zone						,		project=self.gcp_project						)
          else:
             A__         =    tf.distribute.cluster_resolver.TPUClusterResolver()
       except ValueError:
          if self.tpu_name:
             raise RuntimeError(f"""Couldn't connect to TPU {self.tpu_name}!"""						)
          else:
             A__         =    None

       if tpu:
          # Set to bfloat16 in case of TPU
          if self.fpaa:
             tf.keras.mixed_precision.set_global_policy('mixed_bfloat16'						)

          tf.config.experimental_connect_to_cluster(a						)
          tf.tpu.experimental.initialize_tpu_system(a						)

          A__         =    tf.distribute.TPUStrategy(a						)

       elif len(a						) == 0:
          A__         =    tf.distribute.OneDeviceStrategy(device='/cpu:0'						)
       elif len(a						) == 1:
          A__         =    tf.distribute.OneDeviceStrategy(device='/gpu:0'						)
       elif len(a						) > 1:
          # If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`
          A__         =    tf.distribute.MirroredStrategy()
       else:
          raise ValueError('Cannot find the proper strategy, please check your environment properties.'						)

    return strategy






 @property
 def 	_UpperCamelCase   (self					:     List[Any]						)					->      "tf.distribute.Strategy":






    """simple docstring"""


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






 @property
 def 	_UpperCamelCase   (self					:     Tuple						)					->      int:






    """simple docstring"""


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






 @property
 def 	_UpperCamelCase   (self					:     Optional[Any]						)					->      int:






    """simple docstring"""


    return False  # TF Logging is handled by Keras not the Trainer






 @property
 def 	_UpperCamelCase   (self					:     Union[str, Any]						)					->      int:






    """simple docstring"""


    if self.per_gpu_train_batch_size:
       logger.warning(
           'Using deprecated `--per_gpu_train_batch_size` argument which will be removed in a future '
           'version. Using `--per_device_train_batch_size` is preferred.'						)
    A__         =    self.per_gpu_train_batch_size or self.per_device_train_batch_size
    return per_device_batch_size * self.n_replicas






 @property
 def 	_UpperCamelCase   (self					:     Optional[int]						)					->      int:






    """simple docstring"""


    if self.per_gpu_eval_batch_size:
       logger.warning(
           'Using deprecated `--per_gpu_eval_batch_size` argument which will be removed in a future '
           'version. Using `--per_device_eval_batch_size` is preferred.'						)
    A__         =    self.per_gpu_eval_batch_size or self.per_device_eval_batch_size
    return per_device_batch_size * self.n_replicas






 @property
 def 	_UpperCamelCase   (self					:     Optional[Any]						)					->      int:






    """simple docstring"""


    requires_backends(self						,		['tf']						)
    warnings.warn(
        'The n_gpu argument is deprecated and will be removed in a future version, use n_replicas instead.'						,		a						,		)
    return self._setup_strategy.num_replicas_in_sync


