



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 					__lowerCAmelCase :




				def __init__(							self						,				_snake_case						,				_snake_case=13						,				_snake_case=32						,				_snake_case=3						,				_snake_case=4						,				_snake_case=[10, 20, 30, 40]						,				_snake_case=[2, 2, 3, 2]						,				_snake_case=True						,				_snake_case=True						,				_snake_case=37						,				_snake_case="gelu"						,				_snake_case=10						,				_snake_case=0.02						,				_snake_case=["stage2", "stage3", "stage4"]						,				_snake_case=[2, 3, 4]						,				_snake_case=None						,				):



									"""simple docstring"""





									_lowerCAmelCase						=    parent
									_lowerCAmelCase						=    batch_size
									_lowerCAmelCase						=    image_size
									_lowerCAmelCase						=    num_channels
									_lowerCAmelCase						=    num_stages
									_lowerCAmelCase						=    hidden_sizes
									_lowerCAmelCase						=    depths
									_lowerCAmelCase						=    is_training
									_lowerCAmelCase						=    use_labels
									_lowerCAmelCase						=    intermediate_size
									_lowerCAmelCase						=    hidden_act
									_lowerCAmelCase						=    num_labels
									_lowerCAmelCase						=    initializer_range
									_lowerCAmelCase						=    out_features
									_lowerCAmelCase						=    out_indices
									_lowerCAmelCase						=    scope




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

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

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




				def 							snake_case   (							self		):



									"""simple docstring"""





									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=_snake_case						,				initializer_range=self.initializer_range						,				out_features=self.out_features						,				out_indices=self.out_indices						,				num_labels=self.num_labels						,				)




				def 							snake_case   (							self						,				_snake_case						,				_snake_case						,				_snake_case		):



									"""simple docstring"""





									_lowerCAmelCase						=    ConvNextModel(config=_snake_case		)
									model.to(_snake_case		)
									model.eval()
									_lowerCAmelCase						=    model(_snake_case		)
									# 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 							snake_case   (							self						,				_snake_case						,				_snake_case						,				_snake_case		):



									"""simple docstring"""





									_lowerCAmelCase						=    ConvNextForImageClassification(_snake_case		)
									model.to(_snake_case		)
									model.eval()
									_lowerCAmelCase						=    model(_snake_case						,				labels=_snake_case		)
									self.parent.assertEqual(result.logits.shape						,				(self.batch_size, self.num_labels)		)




				def 							snake_case   (							self						,				_snake_case						,				_snake_case						,				_snake_case		):



									"""simple docstring"""





									_lowerCAmelCase						=    ConvNextBackbone(config=_snake_case		)
									model.to(_snake_case		)
									model.eval()
									_lowerCAmelCase						=    model(_snake_case		)

									# 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
									_lowerCAmelCase						=    None
									_lowerCAmelCase						=    ConvNextBackbone(config=_snake_case		)
									model.to(_snake_case		)
									model.eval()
									_lowerCAmelCase						=    model(_snake_case		)

									# 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 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    self.prepare_config_and_inputs()
									_lowerCAmelCase	,	_lowerCAmelCase	,	_lowerCAmelCase						=    config_and_inputs
									_lowerCAmelCase						=    {"""pixel_values""": pixel_values}
									return config, inputs_dict





@require_torch
class 					__lowerCAmelCase (     lowerCamelCase__       ,						lowerCamelCase__       ,						unittest.TestCase       ):
				__lowerCamelCase =				(
				    (
				        ConvNextModel,
				        ConvNextForImageClassification,
				        ConvNextBackbone,
				    )
				    if is_torch_available()
				    else ()
				)
				__lowerCamelCase =				(
				    {'''feature-extraction''': ConvNextModel, '''image-classification''': ConvNextForImageClassification}
				    if is_torch_available()
				    else {}
				)

				__lowerCamelCase =				True
				__lowerCamelCase =				False
				__lowerCamelCase =				False
				__lowerCamelCase =				False
				__lowerCamelCase =				False




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    ConvNextModelTester(self		)
									_lowerCAmelCase						=    ConfigTester(self						,				config_class=_snake_case						,				has_text_modality=_snake_case						,				hidden_size=37		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									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 							snake_case   (							self		):



									"""simple docstring"""





									return




				@unittest.skip(reason="""ConvNext does not use inputs_embeds"""		)
				def 							snake_case   (							self		):



									"""simple docstring"""





									pass




				@unittest.skip(reason="""ConvNext does not support input and output embeddings"""		)
				def 							snake_case   (							self		):



									"""simple docstring"""





									pass




				@unittest.skip(reason="""ConvNext does not use feedforward chunking"""		)
				def 							snake_case   (							self		):



									"""simple docstring"""





									pass




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase	,	_lowerCAmelCase						=    self.model_tester.prepare_config_and_inputs_for_common()

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

														_lowerCAmelCase						=    ["""pixel_values"""]
														self.assertListEqual(arg_names[:1]						,				_snake_case		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*_snake_case		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_backbone(*_snake_case		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									def check_hidden_states_output(_snake_case						,				_snake_case						,				_snake_case		):
														_lowerCAmelCase						=    model_class(_snake_case		)
														model.to(_snake_case		)
														model.eval()

														with torch.no_grad():
																			_lowerCAmelCase						=    model(**self._prepare_for_class(_snake_case						,				_snake_case		)		)

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

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

									_lowerCAmelCase	,	_lowerCAmelCase						=    self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
														_lowerCAmelCase						=    True
														check_hidden_states_output(_snake_case						,				_snake_case						,				_snake_case		)

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

														check_hidden_states_output(_snake_case						,				_snake_case						,				_snake_case		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_classification(*_snake_case		)




				@slow
				def 							snake_case   (							self		):



									"""simple docstring"""





									for model_name in CONVNEXT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														_lowerCAmelCase						=    ConvNextModel.from_pretrained(_snake_case		)
														self.assertIsNotNone(_snake_case		)


def 		_UpperCAmelCase				(      ):





					"""simple docstring"""

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





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




				@cached_property
				def 							snake_case   (							self		):



									"""simple docstring"""





									return AutoImageProcessor.from_pretrained("""facebook/convnext-tiny-224"""		) if is_vision_available() else None




				@slow
				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    ConvNextForImageClassification.from_pretrained("""facebook/convnext-tiny-224"""		).to(_snake_case		)

									_lowerCAmelCase						=    self.default_image_processor
									_lowerCAmelCase						=    prepare_img()
									_lowerCAmelCase						=    image_processor(images=_snake_case						,				return_tensors="""pt"""		).to(_snake_case		)

									# forward pass
									with torch.no_grad():
														_lowerCAmelCase						=    model(**_snake_case		)

									# verify the logits
									_lowerCAmelCase						=    torch.Size((1, 1000)		)
									self.assertEqual(outputs.logits.shape						,				_snake_case		)

									_lowerCAmelCase						=    torch.tensor([-0.0260, -0.4739, 0.1911]		).to(_snake_case		)

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





@require_torch
class 					__lowerCAmelCase (     unittest.TestCase       ,						lowerCamelCase__       ):
				__lowerCamelCase =				(ConvNextBackbone,) if is_torch_available() else ()
				__lowerCamelCase =				ConvNextConfig

				__lowerCamelCase =				False




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    ConvNextModelTester(self		)







import jax.numpy as jnp

from ...utils import logging
from ..ta.modeling_flax_ta import FlaxTaEncoderModel, FlaxTaForConditionalGeneration, FlaxTaModel
from .configuration_mta import MTaConfig


UpperCAmelCase_      : Optional[Any]  					=       logging.get_logger(__name__)

UpperCAmelCase_      : str  					=       'T5Config'







def        SCREAMING_SNAKE_CASE_						(			__A :			jnp.array  ,    __A :			int  ,    __A :			int				)	->		jnp.ndarray:

       """simple docstring"""

       a_					:		Dict							=  jnp.zeros_like(__A				)
       a_					:		Dict							=  shifted_input_ids.at[:, 1:].set(input_ids[:, :-1]				)
       a_					:		str							=  shifted_input_ids.at[:, 0].set(__A				)

       a_					:		int							=  jnp.where(shifted_input_ids == -1_00  ,    __A  ,    __A				)
       return shifted_input_ids
class  SCREAMING_SNAKE_CASE__  (				lowercase__			):
      snake_case__				:     str        =						'''mt5'''
      snake_case__				:     List[Any]        =						MTaConfig
class  SCREAMING_SNAKE_CASE__  (				lowercase__			):
      snake_case__				:     str        =						'''mt5'''
      snake_case__				:     List[str]        =						MTaConfig
class  SCREAMING_SNAKE_CASE__  (				lowercase__			):
      snake_case__				:     Any        =						'''mt5'''
      snake_case__				:     Union[str, Any]        =						MTaConfig


