


from typing import TYPE_CHECKING

from ..utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_tf_available, is_torch_available


snake_case			:      Dict 							=   {
    'configuration_utils': ['GenerationConfig'],
    'streamers': ['TextIteratorStreamer', 'TextStreamer'],
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							snake_case			:      List[Any] 							=   [
							    'Constraint',
							    'ConstraintListState',
							    'DisjunctiveConstraint',
							    'PhrasalConstraint',
							]
							snake_case			:      Tuple 							=   [
							    'BeamHypotheses',
							    'BeamScorer',
							    'BeamSearchScorer',
							    'ConstrainedBeamSearchScorer',
							]
							snake_case			:      str 							=   [
							    'EpsilonLogitsWarper',
							    'EtaLogitsWarper',
							    'ForcedBOSTokenLogitsProcessor',
							    'ForcedEOSTokenLogitsProcessor',
							    'HammingDiversityLogitsProcessor',
							    'InfNanRemoveLogitsProcessor',
							    'LogitsProcessor',
							    'LogitsProcessorList',
							    'LogitsWarper',
							    'MinLengthLogitsProcessor',
							    'MinNewTokensLengthLogitsProcessor',
							    'NoBadWordsLogitsProcessor',
							    'NoRepeatNGramLogitsProcessor',
							    'PrefixConstrainedLogitsProcessor',
							    'RepetitionPenaltyLogitsProcessor',
							    'SequenceBiasLogitsProcessor',
							    'EncoderRepetitionPenaltyLogitsProcessor',
							    'TemperatureLogitsWarper',
							    'TopKLogitsWarper',
							    'TopPLogitsWarper',
							    'TypicalLogitsWarper',
							    'EncoderNoRepeatNGramLogitsProcessor',
							    'ExponentialDecayLengthPenalty',
							    'LogitNormalization',
							]
							snake_case			:      Any 							=   [
							    'MaxNewTokensCriteria',
							    'MaxLengthCriteria',
							    'MaxTimeCriteria',
							    'StoppingCriteria',
							    'StoppingCriteriaList',
							    'validate_stopping_criteria',
							]
							snake_case			:      Optional[Any] 							=   [
							    'GenerationMixin',
							    'top_k_top_p_filtering',
							    'GreedySearchEncoderDecoderOutput',
							    'GreedySearchDecoderOnlyOutput',
							    'SampleEncoderDecoderOutput',
							    'SampleDecoderOnlyOutput',
							    'BeamSearchEncoderDecoderOutput',
							    'BeamSearchDecoderOnlyOutput',
							    'BeamSampleEncoderDecoderOutput',
							    'BeamSampleDecoderOnlyOutput',
							    'ContrastiveSearchEncoderDecoderOutput',
							    'ContrastiveSearchDecoderOnlyOutput',
							]

try:
							if not is_tf_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							snake_case			:      Union[str, Any] 							=   [
							    'TFForcedBOSTokenLogitsProcessor',
							    'TFForcedEOSTokenLogitsProcessor',
							    'TFLogitsProcessor',
							    'TFLogitsProcessorList',
							    'TFLogitsWarper',
							    'TFMinLengthLogitsProcessor',
							    'TFNoBadWordsLogitsProcessor',
							    'TFNoRepeatNGramLogitsProcessor',
							    'TFRepetitionPenaltyLogitsProcessor',
							    'TFTemperatureLogitsWarper',
							    'TFTopKLogitsWarper',
							    'TFTopPLogitsWarper',
							    'TFForceTokensLogitsProcessor',
							    'TFSuppressTokensAtBeginLogitsProcessor',
							    'TFSuppressTokensLogitsProcessor',
							]
							snake_case			:      str 							=   [
							    'TFGenerationMixin',
							    'tf_top_k_top_p_filtering',
							    'TFGreedySearchDecoderOnlyOutput',
							    'TFGreedySearchEncoderDecoderOutput',
							    'TFSampleEncoderDecoderOutput',
							    'TFSampleDecoderOnlyOutput',
							    'TFBeamSearchEncoderDecoderOutput',
							    'TFBeamSearchDecoderOnlyOutput',
							    'TFBeamSampleEncoderDecoderOutput',
							    'TFBeamSampleDecoderOnlyOutput',
							    'TFContrastiveSearchEncoderDecoderOutput',
							    'TFContrastiveSearchDecoderOnlyOutput',
							]

try:
							if not is_flax_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							snake_case			:      Union[str, Any] 							=   [
							    'FlaxForcedBOSTokenLogitsProcessor',
							    'FlaxForcedEOSTokenLogitsProcessor',
							    'FlaxLogitsProcessor',
							    'FlaxLogitsProcessorList',
							    'FlaxLogitsWarper',
							    'FlaxMinLengthLogitsProcessor',
							    'FlaxTemperatureLogitsWarper',
							    'FlaxTopKLogitsWarper',
							    'FlaxTopPLogitsWarper',
							]
							snake_case			:      List[str] 							=   [
							    'FlaxGenerationMixin',
							    'FlaxGreedySearchOutput',
							    'FlaxSampleOutput',
							    'FlaxBeamSearchOutput',
							]

if TYPE_CHECKING:
							from .configuration_utils import GenerationConfig
							from .streamers import TextIteratorStreamer, TextStreamer

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .beam_constraints import Constraint, ConstraintListState, DisjunctiveConstraint, PhrasalConstraint
														from .beam_search import BeamHypotheses, BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer
														from .logits_process import (
														    EncoderNoRepeatNGramLogitsProcessor,
														    EncoderRepetitionPenaltyLogitsProcessor,
														    EpsilonLogitsWarper,
														    EtaLogitsWarper,
														    ExponentialDecayLengthPenalty,
														    ForcedBOSTokenLogitsProcessor,
														    ForcedEOSTokenLogitsProcessor,
														    HammingDiversityLogitsProcessor,
														    InfNanRemoveLogitsProcessor,
														    LogitNormalization,
														    LogitsProcessor,
														    LogitsProcessorList,
														    LogitsWarper,
														    MinLengthLogitsProcessor,
														    MinNewTokensLengthLogitsProcessor,
														    NoBadWordsLogitsProcessor,
														    NoRepeatNGramLogitsProcessor,
														    PrefixConstrainedLogitsProcessor,
														    RepetitionPenaltyLogitsProcessor,
														    SequenceBiasLogitsProcessor,
														    TemperatureLogitsWarper,
														    TopKLogitsWarper,
														    TopPLogitsWarper,
														    TypicalLogitsWarper,
														)
														from .stopping_criteria import (
														    MaxLengthCriteria,
														    MaxNewTokensCriteria,
														    MaxTimeCriteria,
														    StoppingCriteria,
														    StoppingCriteriaList,
														    validate_stopping_criteria,
														)
														from .utils import (
														    BeamSampleDecoderOnlyOutput,
														    BeamSampleEncoderDecoderOutput,
														    BeamSearchDecoderOnlyOutput,
														    BeamSearchEncoderDecoderOutput,
														    ContrastiveSearchDecoderOnlyOutput,
														    ContrastiveSearchEncoderDecoderOutput,
														    GenerationMixin,
														    GreedySearchDecoderOnlyOutput,
														    GreedySearchEncoderDecoderOutput,
														    SampleDecoderOnlyOutput,
														    SampleEncoderDecoderOutput,
														    top_k_top_p_filtering,
														)

							try:
														if not is_tf_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .tf_logits_process import (
														    TFForcedBOSTokenLogitsProcessor,
														    TFForcedEOSTokenLogitsProcessor,
														    TFForceTokensLogitsProcessor,
														    TFLogitsProcessor,
														    TFLogitsProcessorList,
														    TFLogitsWarper,
														    TFMinLengthLogitsProcessor,
														    TFNoBadWordsLogitsProcessor,
														    TFNoRepeatNGramLogitsProcessor,
														    TFRepetitionPenaltyLogitsProcessor,
														    TFSuppressTokensAtBeginLogitsProcessor,
														    TFSuppressTokensLogitsProcessor,
														    TFTemperatureLogitsWarper,
														    TFTopKLogitsWarper,
														    TFTopPLogitsWarper,
														)
														from .tf_utils import (
														    TFBeamSampleDecoderOnlyOutput,
														    TFBeamSampleEncoderDecoderOutput,
														    TFBeamSearchDecoderOnlyOutput,
														    TFBeamSearchEncoderDecoderOutput,
														    TFContrastiveSearchDecoderOnlyOutput,
														    TFContrastiveSearchEncoderDecoderOutput,
														    TFGenerationMixin,
														    TFGreedySearchDecoderOnlyOutput,
														    TFGreedySearchEncoderDecoderOutput,
														    TFSampleDecoderOnlyOutput,
														    TFSampleEncoderDecoderOutput,
														    tf_top_k_top_p_filtering,
														)

							try:
														if not is_flax_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .flax_logits_process import (
														    FlaxForcedBOSTokenLogitsProcessor,
														    FlaxForcedEOSTokenLogitsProcessor,
														    FlaxLogitsProcessor,
														    FlaxLogitsProcessorList,
														    FlaxLogitsWarper,
														    FlaxMinLengthLogitsProcessor,
														    FlaxTemperatureLogitsWarper,
														    FlaxTopKLogitsWarper,
														    FlaxTopPLogitsWarper,
														)
														from .flax_utils import FlaxBeamSearchOutput, FlaxGenerationMixin, FlaxGreedySearchOutput, FlaxSampleOutput
else:
							import sys

							snake_case			:      List[Any] 							=   _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




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 				SCREAMING_SNAKE_CASE (     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    __lowercase  (       torch.optim.Optimizer							):

							"""simple docstring"""


							def __init__(    self				, A_				, A_=True				, A_=None	)->	List[str]:
											_SCREAMING_SNAKE_CASE      = optimizer
											_SCREAMING_SNAKE_CASE      = scaler
											_SCREAMING_SNAKE_CASE      = AcceleratorState()
											_SCREAMING_SNAKE_CASE      = GradientState()
											_SCREAMING_SNAKE_CASE      = device_placement
											_SCREAMING_SNAKE_CASE      = False
											_SCREAMING_SNAKE_CASE      = None

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





							@property
							def 		__magic_name__			(    self	)->	Tuple:
											return self.optimizer.state





							@state.setter
							def 		__magic_name__			(    self				, A_	)->	str:
											_SCREAMING_SNAKE_CASE      = state





							@property
							def 		__magic_name__			(    self	)->	Optional[Any]:
											return self.optimizer.param_groups





							@param_groups.setter
							def 		__magic_name__			(    self				, A_	)->	Tuple:
											_SCREAMING_SNAKE_CASE      = param_groups





							@property
							def 		__magic_name__			(    self	)->	Any:
											return self.optimizer.defaults





							@defaults.setter
							def 		__magic_name__			(    self				, A_	)->	Dict:
											_SCREAMING_SNAKE_CASE      = defaults





							def 		__magic_name__			(    self				, A_	)->	Any:
											self.optimizer.add_param_group(A_	)





							def 		__magic_name__			(    self				, A_	)->	int:
											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 		__magic_name__			(    self	)->	List[str]:
											return self.optimizer.state_dict()





							def 		__magic_name__			(    self				, A_=None	)->	Any:
											if self.gradient_state.sync_gradients:
															_SCREAMING_SNAKE_CASE      = 'set_to_none' in inspect.signature(self.optimizer.zero_grad	).parameters
															if accept_arg:
																			if set_to_none is None:
																							_SCREAMING_SNAKE_CASE      = 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 		__magic_name__			(    self				, A_=None	)->	Optional[int]:
											if self.gradient_state.sync_gradients:
															if self.accelerator_state.distributed_type == DistributedType.TPU:
																			_SCREAMING_SNAKE_CASE      = {'closure': closure} if closure is not None else {}
																			xm.optimizer_step(self.optimizer				, optimizer_args=A_	)
															elif self.scaler is not None:
																			_SCREAMING_SNAKE_CASE      = 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
																							_SCREAMING_SNAKE_CASE      = self.scaler.get_scale()
																							_SCREAMING_SNAKE_CASE      = True
																			self.scaler.step(self.optimizer				, A_	)
																			self.scaler.update()
																			_SCREAMING_SNAKE_CASE      = 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.
																							_SCREAMING_SNAKE_CASE      = scale_after < self._last_scale
																			_SCREAMING_SNAKE_CASE      = scale_after

															else:
																			self.optimizer.step(A_	)





							def 		__magic_name__			(    self				, A_	)->	Dict:
											for param_group in self.optimizer.param_groups:
															_SCREAMING_SNAKE_CASE      = [parameters_map.get(A_				, A_	) for p in param_group['params']]





							@property
							def 		__magic_name__			(    self	)->	str:
											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 		__magic_name__			(    self	)->	int:
											return self._is_overflow





							def __getstate__(    self	)->	Dict:
											return self.__dict__.copy()





							def __setstate__(    self				, A_	)->	List[str]:
											self.__dict__.update(A_	)
