



"""simple docstring"""
import inspect
import os
from argparse import ArgumentParser, Namespace
from importlib import import_module

import huggingface_hub
import numpy as np
from packaging import version

from .. import (
    FEATURE_EXTRACTOR_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    PROCESSOR_MAPPING,
    TOKENIZER_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    AutoImageProcessor,
    AutoProcessor,
    AutoTokenizer,
    is_datasets_available,
    is_tf_available,
    is_torch_available,
)
from ..utils import TF2_WEIGHTS_INDEX_NAME, TF2_WEIGHTS_NAME, logging
from . import BaseTransformersCLICommand


if is_tf_available():
						import tensorflow as tf

						tf.config.experimental.enable_tensor_float_aa_execution(False)

if is_torch_available():
						import torch

if is_datasets_available():
						from datasets import load_dataset


__a           =							5e-5  # larger error tolerance than in our internal tests, to avoid flaky user-facing errors
def   A_    (     _lowercase      ):







						'''simple docstring'''


						return PTtoTFCommand(
						    args.model_name,   args.local_dir,   args.max_error,   args.new_weights,   args.no_pr,   args.push,   args.extra_commit_description,   args.override_model_class,   )

class 							lowerCamelCase    (     _lowerCAmelCase							):






					'''simple docstring'''





					@staticmethod
					def   lowerCAmelCase_							(			snake_case:				ArgumentParser		)  ->      Dict:
											snake_case_					:int		=     parser.add_parser(
											    """pt-to-tf"""		,	help=(
											        """CLI tool to run convert a transformers model from a PyTorch checkpoint to a TensorFlow checkpoint."""
											        """ Can also be used to validate existing weights without opening PRs, with --no-pr."""
											    )		,	)
											train_parser.add_argument(
											    """--model-name"""		,	type=snake_case		,	required=snake_case		,	help="""The model name, including owner/organization, as seen on the hub."""		,	)
											train_parser.add_argument(
											    """--local-dir"""		,	type=snake_case		,	default=""""""		,	help="""Optional local directory of the model repository. Defaults to /tmp/{model_name}"""		,	)
											train_parser.add_argument(
											    """--max-error"""		,	type=snake_case		,	default=snake_case		,	help=(
											        f"""Maximum error tolerance. Defaults to {MAX_ERROR}. This flag should be avoided, use at your own risk."""
											    )		,	)
											train_parser.add_argument(
											    """--new-weights"""		,	action="""store_true"""		,	help="""Optional flag to create new TensorFlow weights, even if they already exist."""		,	)
											train_parser.add_argument(
											    """--no-pr"""		,	action="""store_true"""		,	help="""Optional flag to NOT open a PR with converted weights."""		)
											train_parser.add_argument(
											    """--push"""		,	action="""store_true"""		,	help="""Optional flag to push the weights directly to `main` (requires permissions)"""		,	)
											train_parser.add_argument(
											    """--extra-commit-description"""		,	type=snake_case		,	default=""""""		,	help="""Optional additional commit description to use when opening a PR (e.g. to tag the owner)."""		,	)
											train_parser.add_argument(
											    """--override-model-class"""		,	type=snake_case		,	default=snake_case		,	help="""If you think you know better than the auto-detector, you can specify the model class here. """
											    """Can be either an AutoModel class or a specific model class like BertForSequenceClassification."""		,	)
											train_parser.set_defaults(func=snake_case		)






					@staticmethod
					def   lowerCAmelCase_							(			snake_case:				Union[str, Any]		,	snake_case:				List[str]		)  ->      List[str]:
											snake_case_					:List[Any]		=     set(pt_outputs.keys()		)
											snake_case_					:Optional[Any]		=     set(tf_outputs.keys()		)
											if pt_out_attrs != tf_out_attrs:
																	raise ValueError(
																	    f"""The model outputs have different attributes, aborting. (Pytorch: {pt_out_attrs}, TensorFlow:"""
																	    f""" {tf_out_attrs})"""		)

											# 2. For each output attribute, computes the difference
											def _find_pt_tf_differences(snake_case:				List[Any]		,	snake_case:				List[Any]		,	snake_case:				Optional[int]		,	snake_case:				int=""		):
																	# If the current attribute is a tensor, it is a leaf and we make the comparison. Otherwise, we will dig in
																	# recursivelly, keeping the name of the attribute.
																	if isinstance(snake_case		,	torch.Tensor		):
																							snake_case_					:List[str]		=     np.max(np.abs(pt_out.numpy() - tf_out.numpy()		)		)
																							snake_case_					:Optional[int]		=     tensor_difference
																	else:
																							snake_case_					:Optional[int]		=     attr_name
																							for i, pt_item in enumerate(snake_case		):
																													# If it is a named attribute, we keep the name. Otherwise, just its index.
																													if isinstance(snake_case		,	snake_case		):
																																			snake_case_					:int		=     root_name + pt_item
																																			snake_case_					:Dict		=     tf_out[pt_item]
																																			snake_case_					:List[str]		=     pt_out[pt_item]
																													else:
																																			snake_case_					:str		=     root_name + f"""[{i}]"""
																																			snake_case_					:List[Any]		=     tf_out[i]
																													snake_case_					:Dict		=     _find_pt_tf_differences(snake_case		,	snake_case		,	snake_case		,	snake_case		)

																	return differences

											return _find_pt_tf_differences(snake_case		,	snake_case		,	{}		)






					def __init__(			self:				Tuple		,	snake_case:				str		,	snake_case:				str		,	snake_case:				float		,	snake_case:				bool		,	snake_case:				bool		,	snake_case:				bool		,	snake_case:				str		,	snake_case:				str		,	*snake_case:				Any		,	)  ->      Union[str, Any]:
											snake_case_					:List[Any]		=     logging.get_logger("""transformers-cli/pt_to_tf"""		)
											snake_case_					:Optional[int]		=     model_name
											snake_case_					:List[str]		=     local_dir if local_dir else os.path.join("""/tmp"""		,	snake_case		)
											snake_case_					:Optional[Any]		=     max_error
											snake_case_					:List[str]		=     new_weights
											snake_case_					:List[str]		=     no_pr
											snake_case_					:Dict		=     push
											snake_case_					:Tuple		=     extra_commit_description
											snake_case_					:Optional[Any]		=     override_model_class






					def   lowerCAmelCase_							(			self:				Tuple		,	snake_case:				List[Any]		,	snake_case:				str		,	snake_case:				Optional[Any]		)  ->      Optional[int]:

											def _get_audio_input():
																	snake_case_					:int		=     load_dataset("""hf-internal-testing/librispeech_asr_dummy"""		,	"""clean"""		,	split="""validation"""		)
																	snake_case_					:int		=     ds.sort("""id"""		).select(range(2		)		)[:2]["""audio"""]
																	snake_case_					:int		=     [x["""array"""] for x in speech_samples]
																	return raw_samples

											snake_case_					:List[str]		=     type(pt_model.config		)
											if model_config_class in PROCESSOR_MAPPING:
																	snake_case_					:Any		=     AutoProcessor.from_pretrained(self._local_dir		)
																	if model_config_class in TOKENIZER_MAPPING and processor.tokenizer.pad_token is None:
																							snake_case_					:List[str]		=     processor.tokenizer.eos_token
											elif model_config_class in IMAGE_PROCESSOR_MAPPING:
																	snake_case_					:List[Any]		=     AutoImageProcessor.from_pretrained(self._local_dir		)
											elif model_config_class in FEATURE_EXTRACTOR_MAPPING:
																	snake_case_					:int		=     AutoFeatureExtractor.from_pretrained(self._local_dir		)
											elif model_config_class in TOKENIZER_MAPPING:
																	snake_case_					:Tuple		=     AutoTokenizer.from_pretrained(self._local_dir		)
																	if processor.pad_token is None:
																							snake_case_					:str		=     processor.eos_token
											else:
																	raise ValueError(f"""Unknown data processing type (model config type: {model_config_class})"""		)

											snake_case_					:Optional[int]		=     set(inspect.signature(pt_model.forward		).parameters.keys()		)
											snake_case_					:Dict		=     {}
											if "input_ids" in model_forward_signature:
																	processor_inputs.update(
																	    {
																	        """text""": ["""Hi there!""", """I am a batch with more than one row and different input lengths."""],
																	        """padding""": True,
																	        """truncation""": True,
																	    }		)
											if "pixel_values" in model_forward_signature:
																	snake_case_					:Optional[int]		=     load_dataset("""cifar10"""		,	"""plain_text"""		,	split="""test"""		)[:2]["""img"""]
																	processor_inputs.update({"""images""": sample_images}		)
											if "input_features" in model_forward_signature:
																	snake_case_					:List[str]		=     inspect.signature(processor.feature_extractor		).parameters
																	# Pad to the largest input length by default but take feature extractor default
																	# padding value if it exists e.g. "max_length" and is not False or None
																	if "padding" in feature_extractor_signature:
																							snake_case_					:Tuple		=     feature_extractor_signature["""padding"""].default
																							if default_strategy is not False and default_strategy is not None:
																													snake_case_					:Any		=     default_strategy
																							else:
																													snake_case_					:Optional[Any]		=     True
																	else:
																							snake_case_					:List[Any]		=     True
																	processor_inputs.update({"""audio""": _get_audio_input(), """padding""": padding_strategy}		)
											if "input_values" in model_forward_signature:  # Wav2Vec2 audio input
																	processor_inputs.update({"""audio""": _get_audio_input(), """padding""": True}		)
											snake_case_					:Any		=     processor(**snake_case		,	return_tensors="""pt"""		)
											snake_case_					:Union[str, Any]		=     processor(**snake_case		,	return_tensors="""tf"""		)

											# Extra input requirements, in addition to the input modality
											if (
											    config.is_encoder_decoder
											    or (hasattr(snake_case		,	"""encoder"""		) and hasattr(snake_case		,	"""decoder"""		))
											    or "decoder_input_ids" in tf_dummy_inputs
											):
																	snake_case_					:Optional[int]		=     np.asarray([[1], [1]]		,	dtype=snake_case		) * (pt_model.config.decoder_start_token_id or 0)
																	pt_input.update({"""decoder_input_ids""": torch.tensor(snake_case		)}		)
																	tf_input.update({"""decoder_input_ids""": tf.convert_to_tensor(snake_case		)}		)

											return pt_input, tf_input






					def   lowerCAmelCase_							(			self:				Tuple		)  ->      List[str]:
											# hub version 0.9.0 introduced the possibility of programmatically opening PRs with normal write tokens.
											if version.parse(huggingface_hub.__version__		) < version.parse("""0.9.0"""		):
																	raise ImportError(
																	    """The huggingface_hub version must be >= 0.9.0 to use this command. Please update your huggingface_hub"""
																	    """ installation."""		)
											else:
																	from huggingface_hub import Repository, create_commit
																	from huggingface_hub._commit_api import CommitOperationAdd

											# Fetch remote data
											snake_case_					:Tuple		=     Repository(local_dir=self._local_dir		,	clone_from=self._model_name		)

											# Load config and get the appropriate architecture -- the latter is needed to convert the head's weights
											snake_case_					:Dict		=     AutoConfig.from_pretrained(self._local_dir		)
											snake_case_					:Union[str, Any]		=     config.architectures
											if self._override_model_class is not None:
																	if self._override_model_class.startswith("""TF"""		):
																							snake_case_					:int		=     [self._override_model_class[2:]]
																	else:
																							snake_case_					:Tuple		=     [self._override_model_class]
																	try:
																							snake_case_					:Optional[Any]		=     getattr(import_module("""transformers"""		)		,	architectures[0]		)
																	except AttributeError:
																							raise ValueError(f"""Model class {self._override_model_class} not found in transformers."""		)
																	try:
																							snake_case_					:List[str]		=     getattr(import_module("""transformers"""		)		,	"""TF""" + architectures[0]		)
																	except AttributeError:
																							raise ValueError(f"""TF model class TF{self._override_model_class} not found in transformers."""		)
											elif architectures is None:  # No architecture defined -- use auto classes
																	snake_case_					:Any		=     getattr(import_module("""transformers"""		)		,	"""AutoModel"""		)
																	snake_case_					:int		=     getattr(import_module("""transformers"""		)		,	"""TFAutoModel"""		)
																	self._logger.warning("""No detected architecture, using AutoModel/TFAutoModel"""		)
											else:  # Architecture defined -- use it
																	if len(snake_case		) > 1:
																							raise ValueError(f"""More than one architecture was found, aborting. (architectures = {architectures})"""		)
																	self._logger.warning(f"""Detected architecture: {architectures[0]}"""		)
																	snake_case_					:str		=     getattr(import_module("""transformers"""		)		,	architectures[0]		)
																	try:
																							snake_case_					:Optional[Any]		=     getattr(import_module("""transformers"""		)		,	"""TF""" + architectures[0]		)
																	except AttributeError:
																							raise AttributeError(f"""The TensorFlow equivalent of {architectures[0]} doesn't exist in transformers."""		)

        # Check the TF dummy inputs to see what keys we need in the forward pass
											snake_case_					:List[str]		=     tf_class.from_config(snake_case		)
											snake_case_					:str		=     tf_from_pt_model.dummy_inputs

											del tf_from_pt_model  # Try to keep only one model in memory at a time

											# Load the model and get some basic inputs
											snake_case_					:Any		=     pt_class.from_pretrained(self._local_dir		)
											pt_model.eval()

											snake_case_,			snake_case_					:List[Any]		=     self.get_inputs(snake_case		,	snake_case		,	snake_case		)

											with torch.no_grad():
																	snake_case_					:str		=     pt_model(**snake_case		,	output_hidden_states=snake_case		)
											del pt_model  # will no longer be used, and may have a large memory footprint

											snake_case_					:Any		=     tf_class.from_pretrained(self._local_dir		,	from_pt=snake_case		)
											snake_case_					:Dict		=     tf_from_pt_model(**snake_case		,	output_hidden_states=snake_case		,	training=snake_case		)

											# Confirms that cross loading PT weights into TF worked.
											snake_case_					:Optional[Any]		=     self.find_pt_tf_differences(snake_case		,	snake_case		)
											snake_case_					:int		=     {k: v for k, v in crossload_differences.items() if """hidden""" not in k}
											snake_case_					:List[Any]		=     {k: v for k, v in crossload_differences.items() if """hidden""" in k}
											if len(snake_case		) == 0 and architectures is not None:
																	raise ValueError(
																	    f"""Something went wrong -- the config file has architectures ({architectures}), but no model head"""
																	    """ output was found. All outputs start with 'hidden'"""		)
											snake_case_					:Any		=     max(output_differences.values()		) if output_differences else 0.0
											snake_case_					:List[str]		=     max(hidden_differences.values()		)
											if max_crossload_output_diff > self._max_error or max_crossload_hidden_diff > self._max_error:
																	raise ValueError(
																	    """The cross-loaded TensorFlow model has different outputs, something went wrong!\n"""
																	    + f"""\nList of maximum output differences above the threshold ({self._max_error}):\n"""
																	    + """\n""".join([f"""{k}: {v:.3e}""" for k, v in output_differences.items() if v > self._max_error]		)
																	    + f"""\n\nList of maximum hidden layer differences above the threshold ({self._max_error}):\n"""
																	    + """\n""".join([f"""{k}: {v:.3e}""" for k, v in hidden_differences.items() if v > self._max_error]		)		)

											# Save the weights in a TF format (if needed) and confirms that the results are still good
											snake_case_					:Tuple		=     os.path.join(self._local_dir		,	snake_case		)
											snake_case_					:Union[str, Any]		=     os.path.join(self._local_dir		,	snake_case		)
											if (not os.path.exists(snake_case		) and not os.path.exists(snake_case		)) or self._new_weights:
																	tf_from_pt_model.save_pretrained(self._local_dir		)
											del tf_from_pt_model  # will no longer be used, and may have a large memory footprint

											snake_case_					:Optional[Any]		=     tf_class.from_pretrained(self._local_dir		)
											snake_case_					:Dict		=     tf_model(**snake_case		,	output_hidden_states=snake_case		)

											snake_case_					:List[Any]		=     self.find_pt_tf_differences(snake_case		,	snake_case		)
											snake_case_					:Optional[int]		=     {k: v for k, v in conversion_differences.items() if """hidden""" not in k}
											snake_case_					:List[Any]		=     {k: v for k, v in conversion_differences.items() if """hidden""" in k}
											if len(snake_case		) == 0 and architectures is not None:
																	raise ValueError(
																	    f"""Something went wrong -- the config file has architectures ({architectures}), but no model head"""
																	    """ output was found. All outputs start with 'hidden'"""		)
											snake_case_					:Dict		=     max(output_differences.values()		) if output_differences else 0.0
											snake_case_					:Tuple		=     max(hidden_differences.values()		)
											if max_conversion_output_diff > self._max_error or max_conversion_hidden_diff > self._max_error:
																	raise ValueError(
																	    """The converted TensorFlow model has different outputs, something went wrong!\n"""
																	    + f"""\nList of maximum output differences above the threshold ({self._max_error}):\n"""
																	    + """\n""".join([f"""{k}: {v:.3e}""" for k, v in output_differences.items() if v > self._max_error]		)
																	    + f"""\n\nList of maximum hidden layer differences above the threshold ({self._max_error}):\n"""
																	    + """\n""".join([f"""{k}: {v:.3e}""" for k, v in hidden_differences.items() if v > self._max_error]		)		)

											snake_case_					:Tuple		=     """Update TF weights""" if self._new_weights else """Add TF weights"""
											if self._push:
																	repo.git_add(auto_lfs_track=snake_case		)
																	repo.git_commit(snake_case		)
																	repo.git_push(blocking=snake_case		)  # this prints a progress bar with the upload
																	self._logger.warning(f"""TF weights pushed into {self._model_name}"""		)
											elif not self._no_pr:
																	self._logger.warning("""Uploading the weights into a new PR..."""		)
																	snake_case_					:Tuple		=     (
																	    """Model converted by the [`transformers`' `pt_to_tf`"""
																	    """ CLI](https://github.com/huggingface/transformers/blob/main/src/transformers/commands/pt_to_tf.py). """
																	    """All converted model outputs and hidden layers were validated against its PyTorch counterpart.\n\n"""
																	    f"""Maximum crossload output difference={max_crossload_output_diff:.3e}; """
																	    f"""Maximum crossload hidden layer difference={max_crossload_hidden_diff:.3e};\n"""
																	    f"""Maximum conversion output difference={max_conversion_output_diff:.3e}; """
																	    f"""Maximum conversion hidden layer difference={max_conversion_hidden_diff:.3e};\n"""
																	)
																	if self._max_error > MAX_ERROR:
																							commit_descrition += (
																							    f"""\n\nCAUTION: The maximum admissible error was manually increased to {self._max_error}!"""
																							)
																	if self._extra_commit_description:
																							commit_descrition += "\n\n" + self._extra_commit_description

																	# sharded model -> adds all related files (index and .h5 shards)
																	if os.path.exists(snake_case		):
																							snake_case_					:str		=     [
																							    CommitOperationAdd(path_in_repo=snake_case		,	path_or_fileobj=snake_case		)
																							]
																							for shard_path in tf.io.gfile.glob(self._local_dir + """/tf_model-*.h5"""		):
																													operations += [
																													    CommitOperationAdd(path_in_repo=os.path.basename(snake_case		)		,	path_or_fileobj=snake_case		)
																													]
																	else:
																							snake_case_					:Optional[Any]		=     [CommitOperationAdd(path_in_repo=snake_case		,	path_or_fileobj=snake_case		)]

																	snake_case_					:Tuple		=     create_commit(
																	    repo_id=self._model_name		,	operations=snake_case		,	commit_message=snake_case		,	commit_description=snake_case		,	repo_type="""model"""		,	create_pr=snake_case		,	).pr_url
																	self._logger.warning(f"""PR open in {hub_pr_url}"""		)

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



