



# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation

import warnings

from .state import AcceleratorState, GradientState


warnings.filterwarnings('''ignore''', category=UserWarning, module='''torch.optim.lr_scheduler''')






class 		lowercase						:






	"""simple docstring"""


	def __init__(     self     :						List[str]						,  __UpperCAmelCase     :						Any						,  __UpperCAmelCase     :						List[Any]						,  __UpperCAmelCase     :						bool = True						,  __UpperCAmelCase     :						bool = False					)      ->			int:
							UpperCAmelCase_=       scheduler
							UpperCAmelCase_=       optimizers if isinstance(__UpperCAmelCase						,  (list, tuple)					) else [optimizers]
							UpperCAmelCase_=       split_batches
							UpperCAmelCase_=       step_with_optimizer
							UpperCAmelCase_=       GradientState()




	def   _SCREAMING_SNAKE_CASE		(     self     :						Dict						,  *__UpperCAmelCase     :						str						,  **__UpperCAmelCase     :						Dict					)      ->			List[Any]:
							if not self.step_with_optimizer:
													# No link between scheduler and optimizer -> just step
													self.scheduler.step(*__UpperCAmelCase						,  **__UpperCAmelCase					)
													return

							# Otherwise, first make sure the optimizer was stepped.
							if not self.gradient_state.sync_gradients:
													if self.gradient_state.adjust_scheduler:
																			self.scheduler._step_count += 1
													return

							for opt in self.optimizers:
													if opt.step_was_skipped:
																			return
							if self.split_batches:
													# Split batches -> the training dataloader batch size is not changed so one step per training step
													self.scheduler.step(*__UpperCAmelCase						,  **__UpperCAmelCase					)
							else:
													# Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
													# num_processes steps per training step
													UpperCAmelCase_=       AcceleratorState().num_processes
													for _ in range(__UpperCAmelCase					):
																			# Special case when using OneCycle and `drop_last` was not used
																			if hasattr(self.scheduler						,  """total_steps"""					):
																									if self.scheduler._step_count <= self.scheduler.total_steps:
																															self.scheduler.step(*__UpperCAmelCase						,  **__UpperCAmelCase					)
																			else:
																									self.scheduler.step(*__UpperCAmelCase						,  **__UpperCAmelCase					)




	def   _SCREAMING_SNAKE_CASE		(     self     :						Union[str, Any]					)      ->			List[str]:
							return self.scheduler.get_last_lr()




	def   _SCREAMING_SNAKE_CASE		(     self     :						Optional[Any]					)      ->			Optional[Any]:
							return self.scheduler.state_dict()




	def   _SCREAMING_SNAKE_CASE		(     self     :						Tuple						,  __UpperCAmelCase     :						Any					)      ->			int:
							self.scheduler.load_state_dict(__UpperCAmelCase					)




	def   _SCREAMING_SNAKE_CASE		(     self     :						Tuple					)      ->			List[Any]:
							return self.scheduler.get_lr()




	def   _SCREAMING_SNAKE_CASE		(     self     :						Optional[int]						,  *__UpperCAmelCase     :						Dict						,  **__UpperCAmelCase     :						List[Any]					)      ->			str:
							return self.scheduler.print_lr(*__UpperCAmelCase						,  **__UpperCAmelCase					)





from typing import TYPE_CHECKING

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


__A       =						{
    '''configuration_clip''': [
        '''CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''CLIPConfig''',
        '''CLIPOnnxConfig''',
        '''CLIPTextConfig''',
        '''CLIPVisionConfig''',
    ],
    '''processing_clip''': ['''CLIPProcessor'''],
    '''tokenization_clip''': ['''CLIPTokenizer'''],
}

try:
	if not is_tokenizers_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__A       =						['''CLIPTokenizerFast''']

try:
	if not is_vision_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__A       =						['''CLIPFeatureExtractor''']
	__A       =						['''CLIPImageProcessor''']

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__A       =						[
	    '''CLIP_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''CLIPModel''',
	    '''CLIPPreTrainedModel''',
	    '''CLIPTextModel''',
	    '''CLIPTextModelWithProjection''',
	    '''CLIPVisionModel''',
	    '''CLIPVisionModelWithProjection''',
	]

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__A       =						[
	    '''TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''TFCLIPModel''',
	    '''TFCLIPPreTrainedModel''',
	    '''TFCLIPTextModel''',
	    '''TFCLIPVisionModel''',
	]

try:
	if not is_flax_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__A       =						[
	    '''FlaxCLIPModel''',
	    '''FlaxCLIPPreTrainedModel''',
	    '''FlaxCLIPTextModel''',
	    '''FlaxCLIPTextPreTrainedModel''',
	    '''FlaxCLIPVisionModel''',
	    '''FlaxCLIPVisionPreTrainedModel''',
	]


if TYPE_CHECKING:
	from .configuration_clip import (
	    CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP,
	    CLIPConfig,
	    CLIPOnnxConfig,
	    CLIPTextConfig,
	    CLIPVisionConfig,
	)
	from .processing_clip import CLIPProcessor
	from .tokenization_clip import CLIPTokenizer

	try:
		if not is_tokenizers_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_clip_fast import CLIPTokenizerFast

	try:
		if not is_vision_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .feature_extraction_clip import CLIPFeatureExtractor
		from .image_processing_clip import CLIPImageProcessor

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_clip import (
		    CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
		    CLIPModel,
		    CLIPPreTrainedModel,
		    CLIPTextModel,
		    CLIPTextModelWithProjection,
		    CLIPVisionModel,
		    CLIPVisionModelWithProjection,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_clip import (
		    TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFCLIPModel,
		    TFCLIPPreTrainedModel,
		    TFCLIPTextModel,
		    TFCLIPVisionModel,
		)

	try:
		if not is_flax_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_flax_clip import (
		    FlaxCLIPModel,
		    FlaxCLIPPreTrainedModel,
		    FlaxCLIPTextModel,
		    FlaxCLIPTextPreTrainedModel,
		    FlaxCLIPVisionModel,
		    FlaxCLIPVisionPreTrainedModel,
		)


else:
	import sys

	__A       =						_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)




