



# 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					)









import inspect
import unittest

from transformers import ConvNextConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import ConvNextBackbone, ConvNextForImageClassification, ConvNextModel
	from transformers.models.convnext.modeling_convnext import CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from PIL import Image

	from transformers import AutoImageProcessor






class 		lowercase						:






	"""simple docstring"""


	def __init__(     self     :						Any						,  __UpperCAmelCase     :						Union[str, Any]						,  __UpperCAmelCase     :						Optional[Any]=13						,  __UpperCAmelCase     :						Dict=32						,  __UpperCAmelCase     :						Union[str, Any]=3						,  __UpperCAmelCase     :						int=4						,  __UpperCAmelCase     :						Any=[10, 20, 30, 40]						,  __UpperCAmelCase     :						int=[2, 2, 3, 2]						,  __UpperCAmelCase     :						Tuple=True						,  __UpperCAmelCase     :						Optional[Any]=True						,  __UpperCAmelCase     :						Union[str, Any]=37						,  __UpperCAmelCase     :						List[Any]="gelu"						,  __UpperCAmelCase     :						Dict=10						,  __UpperCAmelCase     :						Dict=0.02						,  __UpperCAmelCase     :						List[str]=["stage2", "stage3", "stage4"]						,  __UpperCAmelCase     :						Dict=[2, 3, 4]						,  __UpperCAmelCase     :						List[str]=None						,  )      ->			List[Any]:
							UpperCAmelCase_=       parent
							UpperCAmelCase_=       batch_size
							UpperCAmelCase_=       image_size
							UpperCAmelCase_=       num_channels
							UpperCAmelCase_=       num_stages
							UpperCAmelCase_=       hidden_sizes
							UpperCAmelCase_=       depths
							UpperCAmelCase_=       is_training
							UpperCAmelCase_=       use_labels
							UpperCAmelCase_=       intermediate_size
							UpperCAmelCase_=       hidden_act
							UpperCAmelCase_=       num_labels
							UpperCAmelCase_=       initializer_range
							UpperCAmelCase_=       out_features
							UpperCAmelCase_=       out_indices
							UpperCAmelCase_=       scope




	def   _SCREAMING_SNAKE_CASE		(     self     :						Any					)      ->			Union[str, Any]:
							UpperCAmelCase_=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

							UpperCAmelCase_=       None
							if self.use_labels:
													UpperCAmelCase_=       ids_tensor([self.batch_size]						,  self.num_labels					)

							UpperCAmelCase_=       self.get_config()
							return config, pixel_values, labels




	def   _SCREAMING_SNAKE_CASE		(     self     :						str					)      ->			List[Any]:
							return ConvNextConfig(
							    num_channels=self.num_channels						,  hidden_sizes=self.hidden_sizes						,  depths=self.depths						,  num_stages=self.num_stages						,  hidden_act=self.hidden_act						,  is_decoder=__UpperCAmelCase						,  initializer_range=self.initializer_range						,  out_features=self.out_features						,  out_indices=self.out_indices						,  num_labels=self.num_labels						,  )




	def   _SCREAMING_SNAKE_CASE		(     self     :						Any						,  __UpperCAmelCase     :						Optional[int]						,  __UpperCAmelCase     :						int						,  __UpperCAmelCase     :						Dict					)      ->			int:
							UpperCAmelCase_=       ConvNextModel(config=__UpperCAmelCase					)
							model.to(__UpperCAmelCase					)
							model.eval()
							UpperCAmelCase_=       model(__UpperCAmelCase					)
							# expected last hidden states: B, C, H // 32, W // 32
							self.parent.assertEqual(
							    result.last_hidden_state.shape						,  (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)						,  )




	def   _SCREAMING_SNAKE_CASE		(     self     :						Union[str, Any]						,  __UpperCAmelCase     :						str						,  __UpperCAmelCase     :						int						,  __UpperCAmelCase     :						Dict					)      ->			List[str]:
							UpperCAmelCase_=       ConvNextForImageClassification(__UpperCAmelCase					)
							model.to(__UpperCAmelCase					)
							model.eval()
							UpperCAmelCase_=       model(__UpperCAmelCase						,  labels=__UpperCAmelCase					)
							self.parent.assertEqual(result.logits.shape						,  (self.batch_size, self.num_labels)					)




	def   _SCREAMING_SNAKE_CASE		(     self     :						int						,  __UpperCAmelCase     :						Tuple						,  __UpperCAmelCase     :						Dict						,  __UpperCAmelCase     :						List[Any]					)      ->			Optional[Any]:
							UpperCAmelCase_=       ConvNextBackbone(config=__UpperCAmelCase					)
							model.to(__UpperCAmelCase					)
							model.eval()
							UpperCAmelCase_=       model(__UpperCAmelCase					)

							# verify hidden states
							self.parent.assertEqual(len(result.feature_maps					)						,  len(config.out_features					)					)
							self.parent.assertListEqual(list(result.feature_maps[0].shape					)						,  [self.batch_size, self.hidden_sizes[1], 4, 4]					)

							# verify channels
							self.parent.assertEqual(len(model.channels					)						,  len(config.out_features					)					)
							self.parent.assertListEqual(model.channels						,  config.hidden_sizes[1:]					)

							# verify backbone works with out_features=None
							UpperCAmelCase_=       None
							UpperCAmelCase_=       ConvNextBackbone(config=__UpperCAmelCase					)
							model.to(__UpperCAmelCase					)
							model.eval()
							UpperCAmelCase_=       model(__UpperCAmelCase					)

							# verify feature maps
							self.parent.assertEqual(len(result.feature_maps					)						,  1					)
							self.parent.assertListEqual(list(result.feature_maps[0].shape					)						,  [self.batch_size, self.hidden_sizes[-1], 1, 1]					)

							# verify channels
							self.parent.assertEqual(len(model.channels					)						,  1					)
							self.parent.assertListEqual(model.channels						,  [config.hidden_sizes[-1]]					)




	def   _SCREAMING_SNAKE_CASE		(     self     :						Optional[Any]					)      ->			List[Any]:
							UpperCAmelCase_=       self.prepare_config_and_inputs()
							UpperCAmelCase_,       UpperCAmelCase_,       UpperCAmelCase_=       config_and_inputs
							UpperCAmelCase_=       {"""pixel_values""": pixel_values}
							return config, inputs_dict






@require_torch
class 		lowercase						(			snake_case__			,   snake_case__			,   unittest.TestCase):






	"""simple docstring"""


	a__				: Tuple		       =							(
	    (
	        ConvNextModel,
	        ConvNextForImageClassification,
	        ConvNextBackbone,
	    )
	    if is_torch_available()
	    else ()
	)
	a__				: Union[str, Any]		       =							(
	    {"feature-extraction": ConvNextModel, "image-classification": ConvNextForImageClassification}
	    if is_torch_available()
	    else {}
	)

	a__				: Optional[int]		       =							True
	a__				: Optional[Any]		       =							False
	a__				: str		       =							False
	a__				: str		       =							False
	a__				: List[str]		       =							False




	def   _SCREAMING_SNAKE_CASE		(     self     :						Optional[Any]					)      ->			Union[str, Any]:
							UpperCAmelCase_=       ConvNextModelTester(self					)
							UpperCAmelCase_=       ConfigTester(self						,  config_class=__UpperCAmelCase						,  has_text_modality=__UpperCAmelCase						,  hidden_size=37					)




	def   _SCREAMING_SNAKE_CASE		(     self     :						Dict					)      ->			Any:
							self.create_and_test_config_common_properties()
							self.config_tester.create_and_test_config_to_json_string()
							self.config_tester.create_and_test_config_to_json_file()
							self.config_tester.create_and_test_config_from_and_save_pretrained()
							self.config_tester.create_and_test_config_with_num_labels()
							self.config_tester.check_config_can_be_init_without_params()
							self.config_tester.check_config_arguments_init()




	def   _SCREAMING_SNAKE_CASE		(     self     :						Tuple					)      ->			Dict:
							return




	@unittest.skip(reason="""ConvNext does not use inputs_embeds"""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						Union[str, Any]					)      ->			List[str]:
							pass




	@unittest.skip(reason="""ConvNext does not support input and output embeddings"""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						Dict					)      ->			Any:
							pass




	@unittest.skip(reason="""ConvNext does not use feedforward chunking"""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						List[Any]					)      ->			Union[str, Any]:
							pass




	def   _SCREAMING_SNAKE_CASE		(     self     :						List[Any]					)      ->			Tuple:
							UpperCAmelCase_,       UpperCAmelCase_=       self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
													UpperCAmelCase_=       model_class(__UpperCAmelCase					)
													UpperCAmelCase_=       inspect.signature(model.forward					)
													# signature.parameters is an OrderedDict => so arg_names order is deterministic
													UpperCAmelCase_=       [*signature.parameters.keys()]

													UpperCAmelCase_=       ["""pixel_values"""]
													self.assertListEqual(arg_names[:1]						,  __UpperCAmelCase					)




	def   _SCREAMING_SNAKE_CASE		(     self     :						Optional[int]					)      ->			str:
							UpperCAmelCase_=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__UpperCAmelCase					)




	def   _SCREAMING_SNAKE_CASE		(     self     :						Dict					)      ->			Dict:
							UpperCAmelCase_=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_backbone(*__UpperCAmelCase					)




	def   _SCREAMING_SNAKE_CASE		(     self     :						Optional[Any]					)      ->			Dict:
							def check_hidden_states_output(__UpperCAmelCase     :						List[Any]						,  __UpperCAmelCase     :						Optional[Any]						,  __UpperCAmelCase     :						Any					):
													UpperCAmelCase_=       model_class(__UpperCAmelCase					)
													model.to(__UpperCAmelCase					)
													model.eval()

													with torch.no_grad():
																			UpperCAmelCase_=       model(**self._prepare_for_class(__UpperCAmelCase						,  __UpperCAmelCase					)					)

													UpperCAmelCase_=       outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

													UpperCAmelCase_=       self.model_tester.num_stages
													self.assertEqual(len(__UpperCAmelCase					)						,  expected_num_stages + 1					)

													# ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
													self.assertListEqual(
													    list(hidden_states[0].shape[-2:]					)						,  [self.model_tester.image_size // 4, self.model_tester.image_size // 4]						,  )

							UpperCAmelCase_,       UpperCAmelCase_=       self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
													UpperCAmelCase_=       True
													check_hidden_states_output(__UpperCAmelCase						,  __UpperCAmelCase						,  __UpperCAmelCase					)

													# check that output_hidden_states also work using config
													del inputs_dict["output_hidden_states"]
													UpperCAmelCase_=       True

													check_hidden_states_output(__UpperCAmelCase						,  __UpperCAmelCase						,  __UpperCAmelCase					)




	def   _SCREAMING_SNAKE_CASE		(     self     :						Tuple					)      ->			Tuple:
							UpperCAmelCase_=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*__UpperCAmelCase					)




	@slow
	def   _SCREAMING_SNAKE_CASE		(     self     :						Dict					)      ->			List[Any]:
							for model_name in CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													UpperCAmelCase_=       ConvNextModel.from_pretrained(__UpperCAmelCase					)
													self.assertIsNotNone(__UpperCAmelCase					)



def 	__a					(   )				->			Optional[Any]:






						'''simple docstring'''



						UpperCAmelCase_=       Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""    )
						return image






@require_torch
@require_vision
class 		lowercase						(			unittest.TestCase):






	"""simple docstring"""


	@cached_property
	def   _SCREAMING_SNAKE_CASE		(     self     :						Dict					)      ->			List[str]:
							return AutoImageProcessor.from_pretrained("""facebook/convnext-tiny-224"""					) if is_vision_available() else None




	@slow
	def   _SCREAMING_SNAKE_CASE		(     self     :						Any					)      ->			List[Any]:
							UpperCAmelCase_=       ConvNextForImageClassification.from_pretrained("""facebook/convnext-tiny-224"""					).to(__UpperCAmelCase					)

							UpperCAmelCase_=       self.default_image_processor
							UpperCAmelCase_=       prepare_img()
							UpperCAmelCase_=       image_processor(images=__UpperCAmelCase						,  return_tensors="""pt"""					).to(__UpperCAmelCase					)

							# forward pass
							with torch.no_grad():
													UpperCAmelCase_=       model(**__UpperCAmelCase					)

							# verify the logits
							UpperCAmelCase_=       torch.Size((1, 1_000)					)
							self.assertEqual(outputs.logits.shape						,  __UpperCAmelCase					)

							UpperCAmelCase_=       torch.tensor([-0.0_260, -0.4_739, 0.1_911]					).to(__UpperCAmelCase					)

							self.assertTrue(torch.allclose(outputs.logits[0, :3]						,  __UpperCAmelCase						,  atol=1E-4					)					)






@require_torch
class 		lowercase						(			unittest.TestCase			,   snake_case__):






	"""simple docstring"""


	a__				: List[str]		       =							(ConvNextBackbone,) if is_torch_available() else ()
	a__				: Dict		       =							ConvNextConfig

	a__				: Union[str, Any]		       =							False




	def   _SCREAMING_SNAKE_CASE		(     self     :						Tuple					)      ->			Union[str, Any]:
							UpperCAmelCase_=       ConvNextModelTester(self					)




