





"""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 	__lowercase  (     snake_case_						:						Optional[int]       ,snake_case_						:						Optional[int]		)     ->Union[str, Any]:



    '''simple docstring'''







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






class 	__snake_case      (	torch.optim.Optimizer  ):





 """simple docstring"""


 def __init__(					self      ,		__lowerCamelCase      ,		__lowerCamelCase=True      ,		__lowerCamelCase=None						):



     '''simple docstring'''




     __A					:		Any						=  optimizer
     __A					:		Tuple						=  scaler
     __A					:		Optional[int]						=  AcceleratorState()
     __A					:		Tuple						=  GradientState()
     __A					:		Tuple						=  device_placement
     __A					:		Any						=  False
     __A					:		List[str]						=  None

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

 @property
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     return self.optimizer.state

 @state.setter
 def       UpperCamelCase__(					self      ,		__lowerCamelCase						):



     '''simple docstring'''




     __A					:		str						=  state

 @property
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     return self.optimizer.param_groups

 @param_groups.setter
 def       UpperCamelCase__(					self      ,		__lowerCamelCase						):



     '''simple docstring'''




     __A					:		Tuple						=  param_groups

 @property
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     return self.optimizer.defaults

 @defaults.setter
 def       UpperCamelCase__(					self      ,		__lowerCamelCase						):



     '''simple docstring'''




     __A					:		List[Any]						=  defaults

 def       UpperCamelCase__(					self      ,		__lowerCamelCase						):



     '''simple docstring'''




     self.optimizer.add_param_group(__lowerCamelCase						)

 def       UpperCamelCase__(					self      ,		__lowerCamelCase						):



     '''simple docstring'''




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

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     return self.optimizer.state_dict()

 def       UpperCamelCase__(					self      ,		__lowerCamelCase=None						):



     '''simple docstring'''




     if self.gradient_state.sync_gradients:
         __A					:		List[str]						=  '''set_to_none''' in inspect.signature(self.optimizer.zero_grad						).parameters
         if accept_arg:
             if set_to_none is None:
                 __A					:		Optional[Any]						=  False
             self.optimizer.zero_grad(set_to_none=__lowerCamelCase						)
         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      ,		__lowerCamelCase=None						):



     '''simple docstring'''




     if self.gradient_state.sync_gradients:
         if self.accelerator_state.distributed_type == DistributedType.TPU:
             __A					:		str						=  {'''closure''': closure} if closure is not None else {}
             xm.optimizer_step(self.optimizer      ,		optimizer_args=__lowerCamelCase						)
         elif self.scaler is not None:
             __A					:		List[str]						=  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					:		str						=  self.scaler.get_scale()
                 __A					:		List[Any]						=  True
             self.scaler.step(self.optimizer      ,		__lowerCamelCase						)
             self.scaler.update()
             __A					:		Dict						=  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					:		List[Any]						=  scale_after < self._last_scale
             __A					:		List[str]						=  scale_after

         else:
             self.optimizer.step(__lowerCamelCase						)

 def       UpperCamelCase__(					self      ,		__lowerCamelCase						):



     '''simple docstring'''




     for param_group in self.optimizer.param_groups:
         __A					:		Any						=  [parameters_map.get(__lowerCamelCase      ,		__lowerCamelCase						) for p in param_group['''params''']]

 @property
 def       UpperCamelCase__(					self						):



     '''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.'''      ,		__lowerCamelCase      ,		)
     return self._is_overflow

 @property
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     return self._is_overflow

 def __getstate__(					self						):



     '''simple docstring'''




     return self.__dict__.copy()







 def __setstate__(					self      ,		__lowerCamelCase						):



     '''simple docstring'''




     self.__dict__.update(__lowerCamelCase						)











"""simple docstring"""



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


a_				=    {
    """configuration_layoutlm""": ["""LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP""", """LayoutLMConfig""", """LayoutLMOnnxConfig"""],
    """tokenization_layoutlm""": ["""LayoutLMTokenizer"""],
}

try:
       if not is_tokenizers_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       a_				=    ["""LayoutLMTokenizerFast"""]

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       a_				=    [
           """LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
           """LayoutLMForMaskedLM""",
           """LayoutLMForSequenceClassification""",
           """LayoutLMForTokenClassification""",
           """LayoutLMForQuestionAnswering""",
           """LayoutLMModel""",
           """LayoutLMPreTrainedModel""",
       ]

try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       a_				=    [
           """TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
           """TFLayoutLMForMaskedLM""",
           """TFLayoutLMForSequenceClassification""",
           """TFLayoutLMForTokenClassification""",
           """TFLayoutLMForQuestionAnswering""",
           """TFLayoutLMMainLayer""",
           """TFLayoutLMModel""",
           """TFLayoutLMPreTrainedModel""",
       ]


if TYPE_CHECKING:
       from .configuration_layoutlm import LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP, LayoutLMConfig, LayoutLMOnnxConfig
       from .tokenization_layoutlm import LayoutLMTokenizer

       try:
              if not is_tokenizers_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .tokenization_layoutlm_fast import LayoutLMTokenizerFast

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_layoutlm import (
                  LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
                  LayoutLMForMaskedLM,
                  LayoutLMForQuestionAnswering,
                  LayoutLMForSequenceClassification,
                  LayoutLMForTokenClassification,
                  LayoutLMModel,
                  LayoutLMPreTrainedModel,
              )
       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_layoutlm import (
                  TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
                  TFLayoutLMForMaskedLM,
                  TFLayoutLMForQuestionAnswering,
                  TFLayoutLMForSequenceClassification,
                  TFLayoutLMForTokenClassification,
                  TFLayoutLMMainLayer,
                  TFLayoutLMModel,
                  TFLayoutLMPreTrainedModel,
              )

else:
       import sys

       a_				=    _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




