






from typing import TYPE_CHECKING

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


snake_case_   :	Union[str, Any]							      ={'''configuration_gptj''': ['''GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''GPTJConfig''', '''GPTJOnnxConfig''']}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					snake_case_   :	int							      =[
					    '''GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST''',
					    '''GPTJForCausalLM''',
					    '''GPTJForQuestionAnswering''',
					    '''GPTJForSequenceClassification''',
					    '''GPTJModel''',
					    '''GPTJPreTrainedModel''',
					]

try:
					if not is_tf_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					snake_case_   :	str							      =[
					    '''TFGPTJForCausalLM''',
					    '''TFGPTJForQuestionAnswering''',
					    '''TFGPTJForSequenceClassification''',
					    '''TFGPTJModel''',
					    '''TFGPTJPreTrainedModel''',
					]

try:
					if not is_flax_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					snake_case_   :	List[Any]							      =[
					    '''FlaxGPTJForCausalLM''',
					    '''FlaxGPTJModel''',
					    '''FlaxGPTJPreTrainedModel''',
					]


if TYPE_CHECKING:
					from .configuration_gptj import GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTJConfig, GPTJOnnxConfig

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_gptj import (
										    GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST,
										    GPTJForCausalLM,
										    GPTJForQuestionAnswering,
										    GPTJForSequenceClassification,
										    GPTJModel,
										    GPTJPreTrainedModel,
										)

					try:
										if not is_tf_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_tf_gptj import (
										    TFGPTJForCausalLM,
										    TFGPTJForQuestionAnswering,
										    TFGPTJForSequenceClassification,
										    TFGPTJModel,
										    TFGPTJPreTrainedModel,
										)

					try:
										if not is_flax_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel, FlaxGPTJPreTrainedModel

else:
					import sys

					snake_case_   :	List[str]							      =_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)


"""simple docstring"""

import collections
import inspect
import unittest

from transformers import NatConfig
from transformers.testing_utils import require_natten, 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, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
				import torch
				from torch import nn

				from transformers import NatBackbone, NatForImageClassification, NatModel
				from transformers.models.nat.modeling_nat import NAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
				from PIL import Image

				from transformers import AutoImageProcessor






class 		_lowercase :



						"""simple docstring"""




						def __init__(				self							:    Tuple     , UpperCamelCase__							:    Any     , UpperCamelCase__							:    Optional[Any]=13     , UpperCamelCase__							:    Optional[Any]=64     , UpperCamelCase__							:    str=4     , UpperCamelCase__							:    List[str]=3     , UpperCamelCase__							:    int=16     , UpperCamelCase__							:    str=[1, 2, 1]     , UpperCamelCase__							:    List[Any]=[2, 4, 8]     , UpperCamelCase__							:    Optional[Any]=3     , UpperCamelCase__							:    List[Any]=2.0     , UpperCamelCase__							:    int=True     , UpperCamelCase__							:    List[str]=0.0     , UpperCamelCase__							:    Optional[Any]=0.0     , UpperCamelCase__							:    Optional[Any]=0.1     , UpperCamelCase__							:    Tuple="gelu"     , UpperCamelCase__							:    Tuple=True     , UpperCamelCase__							:    Any=0.02     , UpperCamelCase__							:    str=1E-5     , UpperCamelCase__							:    Dict=True     , UpperCamelCase__							:    Optional[int]=None     , UpperCamelCase__							:    Any=True     , UpperCamelCase__							:    Tuple=10     , UpperCamelCase__							:    List[Any]=["stage1", "stage2"]     , UpperCamelCase__							:    Optional[int]=[1, 2]     , )				->				Dict:
													'''simple docstring'''



													__UpperCamelCase							=parent
													__UpperCamelCase							=batch_size
													__UpperCamelCase							=image_size
													__UpperCamelCase							=patch_size
													__UpperCamelCase							=num_channels
													__UpperCamelCase							=embed_dim
													__UpperCamelCase							=depths
													__UpperCamelCase							=num_heads
													__UpperCamelCase							=kernel_size
													__UpperCamelCase							=mlp_ratio
													__UpperCamelCase							=qkv_bias
													__UpperCamelCase							=hidden_dropout_prob
													__UpperCamelCase							=attention_probs_dropout_prob
													__UpperCamelCase							=drop_path_rate
													__UpperCamelCase							=hidden_act
													__UpperCamelCase							=patch_norm
													__UpperCamelCase							=layer_norm_eps
													__UpperCamelCase							=initializer_range
													__UpperCamelCase							=is_training
													__UpperCamelCase							=scope
													__UpperCamelCase							=use_labels
													__UpperCamelCase							=num_labels
													__UpperCamelCase							=out_features
													__UpperCamelCase							=out_indices





						def   UpperCAmelCase_ (				self							:    Union[str, Any]				)				->				List[str]:
													'''simple docstring'''



													__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   UpperCAmelCase_ (				self							:    int				)				->				Union[str, Any]:
													'''simple docstring'''



													return NatConfig(
													    num_labels=self.num_labels     , image_size=self.image_size     , patch_size=self.patch_size     , num_channels=self.num_channels     , embed_dim=self.embed_dim     , depths=self.depths     , num_heads=self.num_heads     , kernel_size=self.kernel_size     , mlp_ratio=self.mlp_ratio     , qkv_bias=self.qkv_bias     , hidden_dropout_prob=self.hidden_dropout_prob     , attention_probs_dropout_prob=self.attention_probs_dropout_prob     , drop_path_rate=self.drop_path_rate     , hidden_act=self.hidden_act     , patch_norm=self.patch_norm     , layer_norm_eps=self.layer_norm_eps     , initializer_range=self.initializer_range     , out_features=self.out_features     , out_indices=self.out_indices     , )





						def   UpperCAmelCase_ (				self							:    List[Any]     , UpperCamelCase__							:    Union[str, Any]     , UpperCamelCase__							:    Optional[Any]     , UpperCamelCase__							:    Union[str, Any]				)				->				str:
													'''simple docstring'''



													__UpperCamelCase							=NatModel(config=UpperCamelCase__				)
													model.to(UpperCamelCase__				)
													model.eval()
													__UpperCamelCase							=model(UpperCamelCase__				)

													__UpperCamelCase							=__UpperCamelCase							=(config.image_size // config.patch_size) // (2 ** (len(config.depths				) - 1))
													__UpperCamelCase							=int(config.embed_dim * 2 ** (len(config.depths				) - 1)				)

													self.parent.assertEqual(
													    result.last_hidden_state.shape     , (self.batch_size, expected_height, expected_width, expected_dim)				)





						def   UpperCAmelCase_ (				self							:    Dict     , UpperCamelCase__							:    Union[str, Any]     , UpperCamelCase__							:    List[Any]     , UpperCamelCase__							:    Any				)				->				Dict:
													'''simple docstring'''



													__UpperCamelCase							=NatForImageClassification(UpperCamelCase__				)
													model.to(UpperCamelCase__				)
													model.eval()
													__UpperCamelCase							=model(UpperCamelCase__     , labels=UpperCamelCase__				)
													self.parent.assertEqual(result.logits.shape     , (self.batch_size, self.num_labels)				)

													# test greyscale images
													__UpperCamelCase							=1
													__UpperCamelCase							=NatForImageClassification(UpperCamelCase__				)
													model.to(UpperCamelCase__				)
													model.eval()

													__UpperCamelCase							=floats_tensor([self.batch_size, 1, self.image_size, self.image_size]				)
													__UpperCamelCase							=model(UpperCamelCase__				)
													self.parent.assertEqual(result.logits.shape     , (self.batch_size, self.num_labels)				)





						def   UpperCAmelCase_ (				self							:    Any     , UpperCamelCase__							:    int     , UpperCamelCase__							:    Dict     , UpperCamelCase__							:    str				)				->				List[str]:
													'''simple docstring'''



													__UpperCamelCase							=NatBackbone(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, model.channels[0], 16, 16]				)

													# verify channels
													self.parent.assertEqual(len(model.channels				)     , len(config.out_features				)				)

													# verify backbone works with out_features=None
													__UpperCamelCase							=None
													__UpperCamelCase							=NatBackbone(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, model.channels[-1], 4, 4]				)

													# verify channels
													self.parent.assertEqual(len(model.channels				)     , 1				)





						def   UpperCAmelCase_ (				self							:    Any				)				->				Tuple:
													'''simple docstring'''



													__UpperCamelCase							=self.prepare_config_and_inputs()
													__UpperCamelCase	,      __UpperCamelCase	,      __UpperCamelCase							=config_and_inputs
													__UpperCamelCase							={'''pixel_values''': pixel_values}
													return config, inputs_dict






@require_natten
@require_torch
class 		_lowercase (						__a					,					__a					,					unittest.TestCase      ):



						"""simple docstring"""




						lowercase__       =  (
						    (
						        NatModel,
						        NatForImageClassification,
						        NatBackbone,
						    )
						    if is_torch_available()
						    else ()
						)
						lowercase__       =  (
						    {'''feature-extraction''': NatModel, '''image-classification''': NatForImageClassification}
						    if is_torch_available()
						    else {}
						)
						lowercase__       =  False

						lowercase__       =  False
						lowercase__       =  False
						lowercase__       =  False
						lowercase__       =  False





						def   UpperCAmelCase_ (				self							:    Union[str, Any]				)				->				int:
													'''simple docstring'''



													__UpperCamelCase							=NatModelTester(self				)
													__UpperCamelCase							=ConfigTester(self     , config_class=UpperCamelCase__     , embed_dim=37				)





						def   UpperCAmelCase_ (				self							:    Any				)				->				str:
													'''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   UpperCAmelCase_ (				self							:    Dict				)				->				Optional[int]:
													'''simple docstring'''



													return





						def   UpperCAmelCase_ (				self							:    Union[str, Any]				)				->				Optional[int]:
													'''simple docstring'''



													__UpperCamelCase							=self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*UpperCamelCase__				)





						def   UpperCAmelCase_ (				self							:    List[Any]				)				->				Tuple:
													'''simple docstring'''



													__UpperCamelCase							=self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_image_classification(*UpperCamelCase__				)





						def   UpperCAmelCase_ (				self							:    str				)				->				Union[str, Any]:
													'''simple docstring'''



													__UpperCamelCase							=self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_backbone(*UpperCamelCase__				)





						@unittest.skip(reason='''Nat does not use inputs_embeds'''				)
						def   UpperCAmelCase_ (				self							:    List[str]				)				->				List[str]:
													'''simple docstring'''



													pass





						@unittest.skip(reason='''Nat does not use feedforward chunking'''				)
						def   UpperCAmelCase_ (				self							:    List[Any]				)				->				Optional[Any]:
													'''simple docstring'''



													pass





						def   UpperCAmelCase_ (				self							:    Any				)				->				Optional[Any]:
													'''simple docstring'''



													__UpperCamelCase	,      __UpperCamelCase							=self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				__UpperCamelCase							=model_class(UpperCamelCase__				)
																				self.assertIsInstance(model.get_input_embeddings()     , (nn.Module)				)
																				__UpperCamelCase							=model.get_output_embeddings()
																				self.assertTrue(x is None or isinstance(UpperCamelCase__     , nn.Linear				)				)





						def   UpperCAmelCase_ (				self							:    int				)				->				str:
													'''simple docstring'''



													__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   UpperCAmelCase_ (				self							:    int				)				->				str:
													'''simple docstring'''



													self.skipTest('''Nat\'s attention operation is handled entirely by NATTEN.'''				)





						def   UpperCAmelCase_ (				self							:    Tuple     , UpperCamelCase__							:    Tuple     , UpperCamelCase__							:    Optional[Any]     , UpperCamelCase__							:    List[str]     , UpperCamelCase__							:    int				)				->				str:
													'''simple docstring'''



													__UpperCamelCase							=model_class(UpperCamelCase__				)
													model.to(UpperCamelCase__				)
													model.eval()

													with torch.no_grad():
																				__UpperCamelCase							=model(**self._prepare_for_class(UpperCamelCase__     , UpperCamelCase__				)				)

													__UpperCamelCase							=outputs.hidden_states

													__UpperCamelCase							=getattr(
													    self.model_tester     , '''expected_num_hidden_layers'''     , len(self.model_tester.depths				) + 1				)
													self.assertEqual(len(UpperCamelCase__				)     , UpperCamelCase__				)

													# Nat has a different seq_length
													__UpperCamelCase							=(
													    config.patch_size
													    if isinstance(config.patch_size     , collections.abc.Iterable				)
													    else (config.patch_size, config.patch_size)
													)

													__UpperCamelCase							=image_size[0] // patch_size[0]
													__UpperCamelCase							=image_size[1] // patch_size[1]

													self.assertListEqual(
													    list(hidden_states[0].shape[-3:]				)     , [height, width, self.model_tester.embed_dim]     , )

													if model_class.__name__ != "NatBackbone":
																				__UpperCamelCase							=outputs.reshaped_hidden_states
																				self.assertEqual(len(UpperCamelCase__				)     , UpperCamelCase__				)

																				__UpperCamelCase	,      __UpperCamelCase	,      __UpperCamelCase	,      __UpperCamelCase							=reshaped_hidden_states[0].shape
																				__UpperCamelCase							=(
																				    reshaped_hidden_states[0].view(UpperCamelCase__     , UpperCamelCase__     , UpperCamelCase__     , UpperCamelCase__				).permute(0     , 2     , 3     , 1				)
																				)
																				self.assertListEqual(
																				    list(reshaped_hidden_states.shape[-3:]				)     , [height, width, self.model_tester.embed_dim]     , )





						def   UpperCAmelCase_ (				self							:    Union[str, Any]				)				->				str:
													'''simple docstring'''



													__UpperCamelCase	,      __UpperCamelCase							=self.model_tester.prepare_config_and_inputs_for_common()

													__UpperCamelCase							=(
													    self.model_tester.image_size
													    if isinstance(self.model_tester.image_size     , collections.abc.Iterable				)
													    else (self.model_tester.image_size, self.model_tester.image_size)
													)

													for model_class in self.all_model_classes:
																				__UpperCamelCase							=True
																				self.check_hidden_states_output(UpperCamelCase__     , UpperCamelCase__     , UpperCamelCase__     , UpperCamelCase__				)

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

																				self.check_hidden_states_output(UpperCamelCase__     , UpperCamelCase__     , UpperCamelCase__     , UpperCamelCase__				)





						@slow
						def   UpperCAmelCase_ (				self							:    List[str]				)				->				Dict:
													'''simple docstring'''



													for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																				__UpperCamelCase							=NatModel.from_pretrained(UpperCamelCase__				)
																				self.assertIsNotNone(UpperCamelCase__				)





						def   UpperCAmelCase_ (				self							:    Dict				)				->				int:
													'''simple docstring'''



													__UpperCamelCase	,      __UpperCamelCase							=self.model_tester.prepare_config_and_inputs_for_common()

													__UpperCamelCase							=_config_zero_init(UpperCamelCase__				)
													for model_class in self.all_model_classes:
																				__UpperCamelCase							=model_class(config=UpperCamelCase__				)
																				for name, param in model.named_parameters():
																											if "embeddings" not in name and param.requires_grad:
																																		self.assertIn(
																																		    ((param.data.mean() * 1E9).round() / 1E9).item()     , [0.0, 1.0]     , msg=f"""Parameter {name} of model {model_class} seems not properly initialized"""     , )






@require_natten
@require_vision
@require_torch
class 		_lowercase (						unittest.TestCase      ):



						"""simple docstring"""




						@cached_property
						def   UpperCAmelCase_ (				self							:    Optional[int]				)				->				Any:
													'''simple docstring'''



													return AutoImageProcessor.from_pretrained('''shi-labs/nat-mini-in1k-224'''				) if is_vision_available() else None





						@slow
						def   UpperCAmelCase_ (				self							:    str				)				->				Tuple:
													'''simple docstring'''



													__UpperCamelCase							=NatForImageClassification.from_pretrained('''shi-labs/nat-mini-in1k-224'''				).to(UpperCamelCase__				)
													__UpperCamelCase							=self.default_image_processor

													__UpperCamelCase							=Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''				)
													__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, 1000)				)
													self.assertEqual(outputs.logits.shape     , UpperCamelCase__				)
													__UpperCamelCase							=torch.tensor([0.38_05, -0.86_76, -0.39_12]				).to(UpperCamelCase__				)
													self.assertTrue(torch.allclose(outputs.logits[0, :3]     , UpperCamelCase__     , atol=1E-4				)				)






@require_torch
@require_natten
class 		_lowercase (						unittest.TestCase					,					__a      ):



						"""simple docstring"""




						lowercase__       =  (NatBackbone,) if is_torch_available() else ()
						lowercase__       =  NatConfig





						def   UpperCAmelCase_ (				self							:    Tuple				)				->				Optional[Any]:
													'''simple docstring'''



													__UpperCamelCase							=NatModelTester(self				)



