



'''simple docstring'''

from typing import TYPE_CHECKING

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


UpperCAmelCase__   : List[str]					=		{
    "configuration_openai": ["OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP", "OpenAIGPTConfig"],
    "tokenization_openai": ["OpenAIGPTTokenizer"],
}

try:
		if not is_tokenizers_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		UpperCAmelCase__   : List[str]					=		["OpenAIGPTTokenizerFast"]

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		UpperCAmelCase__   : List[str]					=		[
		    "OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST",
		    "OpenAIGPTDoubleHeadsModel",
		    "OpenAIGPTForSequenceClassification",
		    "OpenAIGPTLMHeadModel",
		    "OpenAIGPTModel",
		    "OpenAIGPTPreTrainedModel",
		    "load_tf_weights_in_openai_gpt",
		]

try:
		if not is_tf_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		UpperCAmelCase__   : List[str]					=		[
		    "TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST",
		    "TFOpenAIGPTDoubleHeadsModel",
		    "TFOpenAIGPTForSequenceClassification",
		    "TFOpenAIGPTLMHeadModel",
		    "TFOpenAIGPTMainLayer",
		    "TFOpenAIGPTModel",
		    "TFOpenAIGPTPreTrainedModel",
		]


if TYPE_CHECKING:
		from .configuration_openai import OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP, OpenAIGPTConfig
		from .tokenization_openai import OpenAIGPTTokenizer

		try:
				if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .tokenization_openai_fast import OpenAIGPTTokenizerFast

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_openai import (
				    OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    OpenAIGPTDoubleHeadsModel,
				    OpenAIGPTForSequenceClassification,
				    OpenAIGPTLMHeadModel,
				    OpenAIGPTModel,
				    OpenAIGPTPreTrainedModel,
				    load_tf_weights_in_openai_gpt,
				)

		try:
				if not is_tf_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_tf_openai import (
				    TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFOpenAIGPTDoubleHeadsModel,
				    TFOpenAIGPTForSequenceClassification,
				    TFOpenAIGPTLMHeadModel,
				    TFOpenAIGPTMainLayer,
				    TFOpenAIGPTModel,
				    TFOpenAIGPTPreTrainedModel,
				)

else:
		import sys

		UpperCAmelCase__   : Any					=		_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)







'''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_					:     Optional[int]						,	UpperCamelCase_					:     str      )					->       List[str]:







			'''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    A     (	torch.optim.Optimizer    ):





	def __init__(					self			:       Union[str, Any]      ,							__magic_name__			:       Optional[int]      ,							__magic_name__			:       List[str]=True      ,							__magic_name__			:       Optional[int]=None							):


				"""simple docstring"""

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

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





	@property
	def 							__SCREAMING_SNAKE_CASE  (					self			:       List[Any]							):


				"""simple docstring"""

				return self.optimizer.state





	@state.setter
	def 							__SCREAMING_SNAKE_CASE  (					self			:       Any      ,							__magic_name__			:       List[Any]							):


				"""simple docstring"""

				lowerCAmelCase__								=							state





	@property
	def 							__SCREAMING_SNAKE_CASE  (					self			:       Any							):


				"""simple docstring"""

				return self.optimizer.param_groups





	@param_groups.setter
	def 							__SCREAMING_SNAKE_CASE  (					self			:       Optional[Any]      ,							__magic_name__			:       Optional[int]							):


				"""simple docstring"""

				lowerCAmelCase__								=							param_groups





	@property
	def 							__SCREAMING_SNAKE_CASE  (					self			:       Optional[int]							):


				"""simple docstring"""

				return self.optimizer.defaults





	@defaults.setter
	def 							__SCREAMING_SNAKE_CASE  (					self			:       Dict      ,							__magic_name__			:       Any							):


				"""simple docstring"""

				lowerCAmelCase__								=							defaults





	def 							__SCREAMING_SNAKE_CASE  (					self			:       Dict      ,							__magic_name__			:       Tuple							):


				"""simple docstring"""

				self.optimizer.add_param_group(__magic_name__							)





	def 							__SCREAMING_SNAKE_CASE  (					self			:       Union[str, Any]      ,							__magic_name__			:       List[Any]							):


				"""simple docstring"""

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





	def 							__SCREAMING_SNAKE_CASE  (					self			:       Tuple							):


				"""simple docstring"""

				return self.optimizer.state_dict()





	def 							__SCREAMING_SNAKE_CASE  (					self			:       Optional[Any]      ,							__magic_name__			:       List[Any]=None							):


				"""simple docstring"""

				if self.gradient_state.sync_gradients:
							lowerCAmelCase__								=							"set_to_none" in inspect.signature(self.optimizer.zero_grad							).parameters
							if accept_arg:
										if set_to_none is None:
													lowerCAmelCase__								=							False
										self.optimizer.zero_grad(set_to_none=__magic_name__							)
							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 							__SCREAMING_SNAKE_CASE  (					self			:       Dict      ,							__magic_name__			:       Dict=None							):


				"""simple docstring"""

				if self.gradient_state.sync_gradients:
							if self.accelerator_state.distributed_type == DistributedType.TPU:
										lowerCAmelCase__								=							{"closure": closure} if closure is not None else {}
										xm.optimizer_step(self.optimizer      ,							optimizer_args=__magic_name__							)
							elif self.scaler is not None:
										lowerCAmelCase__								=							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
													lowerCAmelCase__								=							self.scaler.get_scale()
													lowerCAmelCase__								=							True
										self.scaler.step(self.optimizer      ,							__magic_name__							)
										self.scaler.update()
										lowerCAmelCase__								=							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.
													lowerCAmelCase__								=							scale_after < self._last_scale
										lowerCAmelCase__								=							scale_after

							else:
										self.optimizer.step(__magic_name__							)





	def 							__SCREAMING_SNAKE_CASE  (					self			:       Optional[Any]      ,							__magic_name__			:       List[str]							):


				"""simple docstring"""

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





	@property
	def 							__SCREAMING_SNAKE_CASE  (					self			:       Union[str, Any]							):


				"""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."      ,							__magic_name__      ,							)
				return self._is_overflow





	@property
	def 							__SCREAMING_SNAKE_CASE  (					self			:       List[Any]							):


				"""simple docstring"""

				return self._is_overflow





	def __getstate__(					self			:       int							):


				"""simple docstring"""

				return self.__dict__.copy()







	def __setstate__(					self			:       List[str]      ,							__magic_name__			:       Dict							):


				"""simple docstring"""

				self.__dict__.update(__magic_name__							)





