


from __future__ import annotations

import inspect
import unittest
from typing import List, Tuple

from transformers import ConvNextConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import tensorflow as tf

       from transformers import TFConvNextForImageClassification, TFConvNextModel


if is_vision_available():
       from PIL import Image

       from transformers import ConvNextImageProcessor



class    __lowercase  :

       """simple docstring"""


       def __init__(    self				, A_				, A_=13				, A_=32				, A_=3				, A_=4				, A_=[10, 20, 30, 40]				, A_=[2, 2, 3, 2]				, A_=True				, A_=True				, A_=37				, A_="gelu"				, A_=10				, A_=0.02				, A_=3				, A_=None				, )->	int:
           _SCREAMING_SNAKE_CASE      = parent
           _SCREAMING_SNAKE_CASE      = batch_size
           _SCREAMING_SNAKE_CASE      = image_size
           _SCREAMING_SNAKE_CASE      = num_channels
           _SCREAMING_SNAKE_CASE      = num_stages
           _SCREAMING_SNAKE_CASE      = hidden_sizes
           _SCREAMING_SNAKE_CASE      = depths
           _SCREAMING_SNAKE_CASE      = is_training
           _SCREAMING_SNAKE_CASE      = use_labels
           _SCREAMING_SNAKE_CASE      = intermediate_size
           _SCREAMING_SNAKE_CASE      = hidden_act
           _SCREAMING_SNAKE_CASE      = type_sequence_label_size
           _SCREAMING_SNAKE_CASE      = initializer_range
           _SCREAMING_SNAKE_CASE      = scope





       def 		__magic_name__			(    self	)->	Tuple:
           _SCREAMING_SNAKE_CASE      = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

           _SCREAMING_SNAKE_CASE      = None
           if self.use_labels:
               _SCREAMING_SNAKE_CASE      = ids_tensor([self.batch_size]				, self.type_sequence_label_size	)

           _SCREAMING_SNAKE_CASE      = self.get_config()

           return config, pixel_values, labels





       def 		__magic_name__			(    self	)->	Dict:
           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=A_				, initializer_range=self.initializer_range				, )





       def 		__magic_name__			(    self				, A_				, A_				, A_	)->	Tuple:
           _SCREAMING_SNAKE_CASE      = TFConvNextModel(config=A_	)
           _SCREAMING_SNAKE_CASE      = model(A_				, training=A_	)
           # 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 		__magic_name__			(    self				, A_				, A_				, A_	)->	int:
           _SCREAMING_SNAKE_CASE      = self.type_sequence_label_size
           _SCREAMING_SNAKE_CASE      = TFConvNextForImageClassification(A_	)
           _SCREAMING_SNAKE_CASE      = model(A_				, labels=A_				, training=A_	)
           self.parent.assertEqual(result.logits.shape				, (self.batch_size, self.type_sequence_label_size)	)





       def 		__magic_name__			(    self	)->	Union[str, Any]:
           _SCREAMING_SNAKE_CASE      = self.prepare_config_and_inputs()
           _SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = config_and_inputs
           _SCREAMING_SNAKE_CASE      = {'pixel_values': pixel_values}
           return config, inputs_dict







@require_tf
class    __lowercase  (       UpperCamelCase      ,     UpperCamelCase      ,     unittest.TestCase							):

       """simple docstring"""


       SCREAMING_SNAKE_CASE   :					Any        =						(TFConvNextModel, TFConvNextForImageClassification) if is_tf_available() else ()
       SCREAMING_SNAKE_CASE   :					Any        =						(
           {"feature-extraction": TFConvNextModel, "image-classification": TFConvNextForImageClassification}
           if is_tf_available()
           else {}
       )

       SCREAMING_SNAKE_CASE   :					Optional[Any]        =						False
       SCREAMING_SNAKE_CASE   :					Tuple        =						False
       SCREAMING_SNAKE_CASE   :					Dict        =						False
       SCREAMING_SNAKE_CASE   :					Tuple        =						False
       SCREAMING_SNAKE_CASE   :					Union[str, Any]        =						False





       def 		__magic_name__			(    self	)->	Union[str, Any]:
           _SCREAMING_SNAKE_CASE      = TFConvNextModelTester(self	)
           _SCREAMING_SNAKE_CASE      = ConfigTester(
               self				, config_class=A_				, has_text_modality=A_				, hidden_size=37				, )





       @unittest.skip(reason='ConvNext does not use inputs_embeds'	)
       def 		__magic_name__			(    self	)->	List[Any]:
           pass





       @unittest.skipIf(
           not is_tf_available() or len(tf.config.list_physical_devices('GPU'	)	) == 0				, reason='TF does not support backprop for grouped convolutions on CPU.'				, )
       @slow
       def 		__magic_name__			(    self	)->	int:
           super().test_keras_fit()





       @unittest.skip(reason='ConvNext does not support input and output embeddings'	)
       def 		__magic_name__			(    self	)->	int:
           pass





       def 		__magic_name__			(    self	)->	Union[str, Any]:
           _SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = self.model_tester.prepare_config_and_inputs_for_common()

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

               _SCREAMING_SNAKE_CASE      = ['pixel_values']
               self.assertListEqual(arg_names[:1]				, A_	)





       def 		__magic_name__			(    self	)->	str:
           _SCREAMING_SNAKE_CASE      = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*A_	)





       @unittest.skipIf(
           not is_tf_available() or len(tf.config.list_physical_devices('GPU'	)	) == 0				, reason='TF does not support backprop for grouped convolutions on CPU.'				, )
       def 		__magic_name__			(    self	)->	List[Any]:
           super().test_dataset_conversion()





       def 		__magic_name__			(    self	)->	List[str]:
           def check_hidden_states_output(A_				, A_				, A_	):
               _SCREAMING_SNAKE_CASE      = model_class(A_	)

               _SCREAMING_SNAKE_CASE      = model(**self._prepare_for_class(A_				, A_	)	)
               _SCREAMING_SNAKE_CASE      = outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

               _SCREAMING_SNAKE_CASE      = self.model_tester.num_stages
               self.assertEqual(len(A_	)				, 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]				, )

           _SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               _SCREAMING_SNAKE_CASE      = True
               check_hidden_states_output(A_				, A_				, A_	)

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

               check_hidden_states_output(A_				, A_				, A_	)





       def 		__magic_name__			(    self	)->	str:
           _SCREAMING_SNAKE_CASE      ,      _SCREAMING_SNAKE_CASE      = self.model_tester.prepare_config_and_inputs_for_common()

           def check_equivalence(A_				, A_				, A_				, A_={}	):
               _SCREAMING_SNAKE_CASE      = model(A_				, return_dict=A_				, **A_	)
               _SCREAMING_SNAKE_CASE      = model(A_				, return_dict=A_				, **A_	).to_tuple()

               def recursive_check(A_				, A_	):
                   if isinstance(A_				, (List, Tuple)	):
                       for tuple_iterable_value, dict_iterable_value in zip(A_				, A_	):
                           recursive_check(A_				, A_	)
                   elif tuple_object is None:
                       return
                   else:
                       self.assertTrue(
                           all(tf.equal(A_				, A_	)	)				, msg=(
                               'Tuple and dict output are not equal. Difference:'
                               F''' {tf.math.reduce_max(tf.abs(tuple_object - dict_object	)	)}'''
                           )				, )

                   recursive_check(A_				, A_	)

           for model_class in self.all_model_classes:
               _SCREAMING_SNAKE_CASE      = model_class(A_	)

               _SCREAMING_SNAKE_CASE      = self._prepare_for_class(A_				, A_	)
               _SCREAMING_SNAKE_CASE      = self._prepare_for_class(A_				, A_	)
               check_equivalence(A_				, A_				, A_	)

               _SCREAMING_SNAKE_CASE      = self._prepare_for_class(A_				, A_				, return_labels=A_	)
               _SCREAMING_SNAKE_CASE      = self._prepare_for_class(A_				, A_				, return_labels=A_	)
               check_equivalence(A_				, A_				, A_	)

               _SCREAMING_SNAKE_CASE      = self._prepare_for_class(A_				, A_	)
               _SCREAMING_SNAKE_CASE      = self._prepare_for_class(A_				, A_	)
               check_equivalence(A_				, A_				, A_				, {'output_hidden_states': True}	)

               _SCREAMING_SNAKE_CASE      = self._prepare_for_class(A_				, A_				, return_labels=A_	)
               _SCREAMING_SNAKE_CASE      = self._prepare_for_class(A_				, A_				, return_labels=A_	)
               check_equivalence(A_				, A_				, A_				, {'output_hidden_states': True}	)





       def 		__magic_name__			(    self	)->	Union[str, Any]:
           _SCREAMING_SNAKE_CASE      = self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*A_	)





       @slow
       def 		__magic_name__			(    self	)->	Optional[Any]:
           _SCREAMING_SNAKE_CASE      = TFConvNextModel.from_pretrained('facebook/convnext-tiny-224'	)
           self.assertIsNotNone(A_	)







def 				SCREAMING_SNAKE_CASE (     ):

    """simple docstring"""






    _SCREAMING_SNAKE_CASE      = Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'							)
    return image



@require_tf
@require_vision
class    __lowercase  (       unittest.TestCase							):

       """simple docstring"""


       @cached_property
       def 		__magic_name__			(    self	)->	Dict:
           return ConvNextImageProcessor.from_pretrained('facebook/convnext-tiny-224'	) if is_vision_available() else None





       @slow
       def 		__magic_name__			(    self	)->	Optional[int]:
           _SCREAMING_SNAKE_CASE      = TFConvNextForImageClassification.from_pretrained('facebook/convnext-tiny-224'	)

           _SCREAMING_SNAKE_CASE      = self.default_image_processor
           _SCREAMING_SNAKE_CASE      = prepare_img()
           _SCREAMING_SNAKE_CASE      = image_processor(images=A_				, return_tensors='tf'	)

           # forward pass
           _SCREAMING_SNAKE_CASE      = model(**A_	)

           # verify the logits
           _SCREAMING_SNAKE_CASE      = tf.TensorShape((1, 1000)	)
           self.assertEqual(outputs.logits.shape				, A_	)

           _SCREAMING_SNAKE_CASE      = tf.constant([-0.0260, -0.4739, 0.1911]	)

           tf.debugging.assert_near(outputs.logits[0, :3]				, A_				, atol=1e-4	)




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