




# 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    SCREAMING_SNAKE_CASE_				:



	"""simple docstring"""


	def __init__(					self		:						Tuple					,     lowerCAmelCase		:						Optional[Any]					,     lowerCAmelCase		:						Optional[Any]					,     lowerCAmelCase		:						bool = True					,     lowerCAmelCase		:						bool = False	)	->       str:







						"""simple docstring"""







						__UpperCamelCase					:	str             =      scheduler
						__UpperCamelCase					:	Tuple             =      optimizers if isinstance(lowerCAmelCase					,     (list, tuple)	) else [optimizers]
						__UpperCamelCase					:	Union[str, Any]             =      split_batches
						__UpperCamelCase					:	Any             =      step_with_optimizer
						__UpperCamelCase					:	List[Any]             =      GradientState()






	def 					lowerCamelCase__     (					self		:						Union[str, Any]					,     *lowerCAmelCase		:						List[Any]					,     **lowerCAmelCase		:						str	)	->       str:







						"""simple docstring"""







						if not self.step_with_optimizer:
											# No link between scheduler and optimizer -> just step
											self.scheduler.step(*lowerCAmelCase					,     **lowerCAmelCase	)
											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(*lowerCAmelCase					,     **lowerCAmelCase	)
						else:
											# Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
											# num_processes steps per training step
											__UpperCamelCase					:	Any             =      AcceleratorState().num_processes
											for _ in range(lowerCAmelCase	):
																# 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(*lowerCAmelCase					,     **lowerCAmelCase	)
																else:
																					self.scheduler.step(*lowerCAmelCase					,     **lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						Optional[Any]	)	->       str:







						"""simple docstring"""







						return self.scheduler.get_last_lr()






	def 					lowerCamelCase__     (					self		:						Dict	)	->       Tuple:







						"""simple docstring"""







						return self.scheduler.state_dict()






	def 					lowerCamelCase__     (					self		:						Optional[int]					,     lowerCAmelCase		:						List[str]	)	->       List[str]:







						"""simple docstring"""







						self.scheduler.load_state_dict(lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						Optional[Any]	)	->       List[str]:







						"""simple docstring"""







						return self.scheduler.get_lr()






	def 					lowerCamelCase__     (					self		:						Optional[Any]					,     *lowerCAmelCase		:						int					,     **lowerCAmelCase		:						str	)	->       int:







						"""simple docstring"""







						return self.scheduler.print_lr(*lowerCAmelCase					,     **lowerCAmelCase	)












# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


a__		     =     {
    '''configuration_efficientnet''': [
        '''EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''EfficientNetConfig''',
        '''EfficientNetOnnxConfig''',
    ]
}

try:
	if not is_vision_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	a__		     =     ['''EfficientNetImageProcessor''']

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	a__		     =     [
	    '''EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''EfficientNetForImageClassification''',
	    '''EfficientNetModel''',
	    '''EfficientNetPreTrainedModel''',
	]

if TYPE_CHECKING:
	from .configuration_efficientnet import (
	    EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP,
	    EfficientNetConfig,
	    EfficientNetOnnxConfig,
	)

	try:
		if not is_vision_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .image_processing_efficientnet import EfficientNetImageProcessor

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_efficientnet import (
		    EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST,
		    EfficientNetForImageClassification,
		    EfficientNetModel,
		    EfficientNetPreTrainedModel,
		)

else:
	import sys

	a__		     =     _LazyModule(__name__, globals()['''__file__'''], _import_structure)






