






'''simple docstring'''







import os
from pickle import UnpicklingError
from typing import Dict, Tuple

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict, unflatten_dict

import transformers

from .utils import logging


UpperCamelCase     :   int	      =		logging.get_logger(__name__)
def     A__      (					__lowerCAmelCase		:      int ,     __lowerCAmelCase		:      Union[str, Any] ,     __lowerCAmelCase		:      List[Any] ,     __lowerCAmelCase		:      Optional[Any]=False    ):
						try:
												import torch  # noqa: F401
						except ImportError:
												logger.error(
												    """Loading a PyTorch model in Flax, requires both PyTorch and Flax to be installed. Please see"""
												    """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
												    """ instructions."""    )
												raise

						if not is_sharded:
												lowerCamelCase__								=	os.path.abspath(__lowerCAmelCase    )
												logger.info(F'''Loading PyTorch weights from {pt_path}'''    )

												lowerCamelCase__								=	torch.load(__lowerCAmelCase ,     map_location="""cpu"""    )
												logger.info(F'''PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()    ):,} parameters.'''    )

												lowerCamelCase__								=	convert_pytorch_state_dict_to_flax(__lowerCAmelCase ,     __lowerCAmelCase    )
						else:
												# model is sharded and pytorch_checkpoint_path already contains the list of .pt shard files
												lowerCamelCase__								=	convert_pytorch_sharded_state_dict_to_flax(__lowerCAmelCase ,     __lowerCAmelCase    )
						return flax_state_dict
def     A__      (					__lowerCAmelCase		:      Tuple[str] ,     __lowerCAmelCase		:      np.ndarray ,     __lowerCAmelCase		:      Dict[str, jnp.ndarray] ,     __lowerCAmelCase		:      str ,     ):

						def is_key_or_prefix_key_in_dict(__lowerCAmelCase		:      Tuple[str]    ) -> bool:
												return len(set(__lowerCAmelCase    ) & {key, (model_prefix,) + key}    ) > 0

						# layer norm
						lowerCamelCase__								=	pt_tuple_key[:-1] + ("""scale""",)
						if pt_tuple_key[-1] in ["weight", "gamma"] and is_key_or_prefix_key_in_dict(__lowerCAmelCase    ):
												return renamed_pt_tuple_key, pt_tensor

						# batch norm layer mean
						lowerCamelCase__								=	pt_tuple_key[:-1] + ("""mean""",)
						if pt_tuple_key[-1] == "running_mean" and not is_key_or_prefix_key_in_dict(__lowerCAmelCase    ):
												return renamed_pt_tuple_key, pt_tensor

						# batch norm layer var
						lowerCamelCase__								=	pt_tuple_key[:-1] + ("""var""",)
						if pt_tuple_key[-1] == "running_var" and not is_key_or_prefix_key_in_dict(__lowerCAmelCase    ):
												return renamed_pt_tuple_key, pt_tensor

						# embedding
						lowerCamelCase__								=	pt_tuple_key[:-1] + ("""embedding""",)
						if pt_tuple_key[-1] == "weight" and is_key_or_prefix_key_in_dict(__lowerCAmelCase    ):
												return renamed_pt_tuple_key, pt_tensor

						# conv layer
						lowerCamelCase__								=	pt_tuple_key[:-1] + ("""kernel""",)
						if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4 and not is_key_or_prefix_key_in_dict(__lowerCAmelCase    ):
												lowerCamelCase__								=	pt_tensor.transpose(2 ,     3 ,     1 ,     0    )
												return renamed_pt_tuple_key, pt_tensor

						# linear layer
						lowerCamelCase__								=	pt_tuple_key[:-1] + ("""kernel""",)
						if pt_tuple_key[-1] == "weight" and not is_key_or_prefix_key_in_dict(__lowerCAmelCase    ):
												lowerCamelCase__								=	pt_tensor.T
												return renamed_pt_tuple_key, pt_tensor

						# old PyTorch layer norm weight
						lowerCamelCase__								=	pt_tuple_key[:-1] + ("""weight""",)
						if pt_tuple_key[-1] == "gamma":
												return renamed_pt_tuple_key, pt_tensor

						# old PyTorch layer norm bias
						lowerCamelCase__								=	pt_tuple_key[:-1] + ("""bias""",)
						if pt_tuple_key[-1] == "beta":
												return renamed_pt_tuple_key, pt_tensor

						# New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
						lowerCamelCase__								=	None
						if pt_tuple_key[-3::2] == ("parametrizations", "original0"):
												lowerCamelCase__								=	pt_tuple_key[-2] + """_g"""
						elif pt_tuple_key[-3::2] == ("parametrizations", "original1"):
												lowerCamelCase__								=	pt_tuple_key[-2] + """_v"""
						if name is not None:
												lowerCamelCase__								=	pt_tuple_key[:-3] + (name,)
												return renamed_pt_tuple_key, pt_tensor

						return pt_tuple_key, pt_tensor
def     A__      (					__lowerCAmelCase		:      Tuple ,     __lowerCAmelCase		:      Any    ):
						# convert pytorch tensor to numpy
						lowerCamelCase__								=	{k: v.numpy() for k, v in pt_state_dict.items()}

						lowerCamelCase__								=	flax_model.base_model_prefix

						# use params dict if the model contains batch norm layers
						if "params" in flax_model.params:
												lowerCamelCase__								=	flax_model.params["""params"""]
						else:
												lowerCamelCase__								=	flax_model.params
						lowerCamelCase__								=	flatten_dict(__lowerCAmelCase    )

						# add batch_stats keys,values to dict
						if "batch_stats" in flax_model.params:
												lowerCamelCase__								=	flatten_dict(flax_model.params["""batch_stats"""]    )
												random_flax_state_dict.update(__lowerCAmelCase    )

						lowerCamelCase__								=	{}

						lowerCamelCase__								=	(model_prefix not in flax_model_params) and (
						    model_prefix in {k.split("""."""    )[0] for k in pt_state_dict.keys()}
						)
						lowerCamelCase__								=	(model_prefix in flax_model_params) and (
						    model_prefix not in {k.split("""."""    )[0] for k in pt_state_dict.keys()}
						)

						# Need to change some parameters name to match Flax names
						for pt_key, pt_tensor in pt_state_dict.items():
												lowerCamelCase__								=	tuple(pt_key.split("""."""    )    )

												# remove base model prefix if necessary
												lowerCamelCase__								=	pt_tuple_key[0] == model_prefix
												if load_model_with_head_into_base_model and has_base_model_prefix:
																		lowerCamelCase__								=	pt_tuple_key[1:]

												# Correctly rename weight parameters
												lowerCamelCase__							,      lowerCamelCase__								=	rename_key_and_reshape_tensor(
												    __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase    )

												# add model prefix if necessary
												lowerCamelCase__								=	(model_prefix,) + flax_key in random_flax_state_dict
												if load_base_model_into_model_with_head and require_base_model_prefix:
																		lowerCamelCase__								=	(model_prefix,) + flax_key

												if flax_key in random_flax_state_dict:
																		if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
																								raise ValueError(
																								    F'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
																								    F'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'''    )

        # add batch stats if the model contains batchnorm layers
												if "batch_stats" in flax_model.params:
																		if "mean" in flax_key[-1] or "var" in flax_key[-1]:
																								lowerCamelCase__								=	jnp.asarray(__lowerCAmelCase    )
																								continue
																		# remove num_batches_tracked key
																		if "num_batches_tracked" in flax_key[-1]:
																								flax_state_dict.pop(__lowerCAmelCase ,     __lowerCAmelCase    )
																								continue

																		# also add unexpected weight so that warning is thrown
																		lowerCamelCase__								=	jnp.asarray(__lowerCAmelCase    )

												else:
																		# also add unexpected weight so that warning is thrown
																		lowerCamelCase__								=	jnp.asarray(__lowerCAmelCase    )

						return unflatten_dict(__lowerCAmelCase    )
def     A__      (					__lowerCAmelCase		:      Optional[int] ,     __lowerCAmelCase		:      List[str]    ):
						import torch

						# Load the index
						lowerCamelCase__								=	{}
						for shard_file in shard_filenames:
												# load using msgpack utils
												lowerCamelCase__								=	torch.load(__lowerCAmelCase    )
												lowerCamelCase__								=	{k: v.numpy() for k, v in pt_state_dict.items()}

												lowerCamelCase__								=	flax_model.base_model_prefix

												# use params dict if the model contains batch norm layers and then add batch_stats keys,values to dict
												if "batch_stats" in flax_model.params:
																		lowerCamelCase__								=	flax_model.params["""params"""]

																		lowerCamelCase__								=	flatten_dict(__lowerCAmelCase    )
																		random_flax_state_dict.update(flatten_dict(flax_model.params["""batch_stats"""]    )    )
												else:
																		lowerCamelCase__								=	flax_model.params
																		lowerCamelCase__								=	flatten_dict(__lowerCAmelCase    )

												lowerCamelCase__								=	(model_prefix not in flax_model_params) and (
												    model_prefix in {k.split("""."""    )[0] for k in pt_state_dict.keys()}
												)
												lowerCamelCase__								=	(model_prefix in flax_model_params) and (
												    model_prefix not in {k.split("""."""    )[0] for k in pt_state_dict.keys()}
												)
												# Need to change some parameters name to match Flax names
												for pt_key, pt_tensor in pt_state_dict.items():
																		lowerCamelCase__								=	tuple(pt_key.split("""."""    )    )

																		# remove base model prefix if necessary
																		lowerCamelCase__								=	pt_tuple_key[0] == model_prefix
																		if load_model_with_head_into_base_model and has_base_model_prefix:
																								lowerCamelCase__								=	pt_tuple_key[1:]

																		# Correctly rename weight parameters
																		lowerCamelCase__							,      lowerCamelCase__								=	rename_key_and_reshape_tensor(
																		    __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase    )
																		# add model prefix if necessary
																		lowerCamelCase__								=	(model_prefix,) + flax_key in random_flax_state_dict
																		if load_base_model_into_model_with_head and require_base_model_prefix:
																								lowerCamelCase__								=	(model_prefix,) + flax_key

																		if flax_key in random_flax_state_dict:
																								if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
																														raise ValueError(
																														    F'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
																														    F'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'''    )

            # add batch stats if the model contains batchnorm layers
																		if "batch_stats" in flax_model.params:
																								if "mean" in flax_key[-1]:
																														lowerCamelCase__								=	jnp.asarray(__lowerCAmelCase    )
																														continue
																								if "var" in flax_key[-1]:
																														lowerCamelCase__								=	jnp.asarray(__lowerCAmelCase    )
																														continue
																								# remove num_batches_tracked key
																								if "num_batches_tracked" in flax_key[-1]:
																														flax_state_dict.pop(__lowerCAmelCase ,     __lowerCAmelCase    )
																														continue

																								# also add unexpected weight so that warning is thrown
																								lowerCamelCase__								=	jnp.asarray(__lowerCAmelCase    )

																		else:
																								# also add unexpected weight so that warning is thrown
																								lowerCamelCase__								=	jnp.asarray(__lowerCAmelCase    )
						return unflatten_dict(__lowerCAmelCase    )
def     A__      (					__lowerCAmelCase		:      Any ,     __lowerCAmelCase		:      Tuple    ):
						lowerCamelCase__								=	os.path.abspath(__lowerCAmelCase    )
						logger.info(F'''Loading Flax weights from {flax_checkpoint_path}'''    )

						# import correct flax class
						lowerCamelCase__								=	getattr(__lowerCAmelCase ,     """Flax""" + model.__class__.__name__    )

						# load flax weight dict
						with open(__lowerCAmelCase ,     """rb"""    ) as state_f:
												try:
																		lowerCamelCase__								=	from_bytes(__lowerCAmelCase ,     state_f.read()    )
												except UnpicklingError:
																		raise EnvironmentError(F'''Unable to convert {flax_checkpoint_path} to Flax deserializable object. '''    )

						return load_flax_weights_in_pytorch_model(__lowerCAmelCase ,     __lowerCAmelCase    )





def     A__      (					__lowerCAmelCase		:      str ,     __lowerCAmelCase		:      Tuple    ):

						try:
												import torch  # noqa: F401
						except ImportError:
												logger.error(
												    """Loading a Flax weights in PyTorch, requires both PyTorch and Flax to be installed. Please see"""
												    """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
												    """ instructions."""    )
												raise

						# check if we have bf16 weights
						lowerCamelCase__								=	flatten_dict(jax.tree_util.tree_map(lambda __lowerCAmelCase    : x.dtype == jnp.bfloataa ,     __lowerCAmelCase    )    ).values()
						if any(__lowerCAmelCase    ):
												# convert all weights to fp32 if the are bf16 since torch.from_numpy can-not handle bf16
												# and bf16 is not fully supported in PT yet.
												logger.warning(
												    """Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` """
												    """before loading those in PyTorch model."""    )
												lowerCamelCase__								=	jax.tree_util.tree_map(
												    lambda __lowerCAmelCase    : params.astype(np.floataa    ) if params.dtype == jnp.bfloataa else params ,     __lowerCAmelCase    )

						lowerCamelCase__								=	flatten_dict(__lowerCAmelCase    )
						lowerCamelCase__								=	pt_model.state_dict()

						lowerCamelCase__								=	(pt_model.base_model_prefix in flax_state) and (
						    pt_model.base_model_prefix not in {k.split("""."""    )[0] for k in pt_model_dict.keys()}
						)
						lowerCamelCase__								=	(pt_model.base_model_prefix not in flax_state) and (
						    pt_model.base_model_prefix in {k.split("""."""    )[0] for k in pt_model_dict.keys()}
						)

						# keep track of unexpected & missing keys
						lowerCamelCase__								=	[]
						lowerCamelCase__								=	set(pt_model_dict.keys()    )

						for flax_key_tuple, flax_tensor in flax_state_dict.items():
												lowerCamelCase__								=	flax_key_tuple[0] == pt_model.base_model_prefix
												lowerCamelCase__								=	""".""".join((pt_model.base_model_prefix,) + flax_key_tuple    ) in pt_model_dict

												# adapt flax_key to prepare for loading from/to base model only
												if load_model_with_head_into_base_model and has_base_model_prefix:
																		lowerCamelCase__								=	flax_key_tuple[1:]
												elif load_base_model_into_model_with_head and require_base_model_prefix:
																		lowerCamelCase__								=	(pt_model.base_model_prefix,) + flax_key_tuple

												# rename flax weights to PyTorch format
												if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 4 and ".".join(__lowerCAmelCase    ) not in pt_model_dict:
																		# conv layer
																		lowerCamelCase__								=	flax_key_tuple[:-1] + ("""weight""",)
																		lowerCamelCase__								=	jnp.transpose(__lowerCAmelCase ,     (3, 2, 0, 1)    )
												elif flax_key_tuple[-1] == "kernel" and ".".join(__lowerCAmelCase    ) not in pt_model_dict:
																		# linear layer
																		lowerCamelCase__								=	flax_key_tuple[:-1] + ("""weight""",)
																		lowerCamelCase__								=	flax_tensor.T
												elif flax_key_tuple[-1] in ["scale", "embedding"]:
																		lowerCamelCase__								=	flax_key_tuple[:-1] + ("""weight""",)

												# adding batch stats from flax batch norm to pt
												elif "mean" in flax_key_tuple[-1]:
																		lowerCamelCase__								=	flax_key_tuple[:-1] + ("""running_mean""",)
												elif "var" in flax_key_tuple[-1]:
																		lowerCamelCase__								=	flax_key_tuple[:-1] + ("""running_var""",)

												if "batch_stats" in flax_state:
																		lowerCamelCase__								=	""".""".join(flax_key_tuple[1:]    )  # Remove the params/batch_stats header
												else:
																		lowerCamelCase__								=	""".""".join(__lowerCAmelCase    )

												# We also need to look at `pt_model_dict` and see if there are keys requiring further transformation.
												lowerCamelCase__								=	{}
												# New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
												for key in pt_model_dict:
																		lowerCamelCase__								=	key.split("""."""    )
																		lowerCamelCase__								=	None
																		if key_components[-3::2] == ["parametrizations", "original0"]:
																								lowerCamelCase__								=	key_components[-2] + """_g"""
																		elif key_components[-3::2] == ["parametrizations", "original1"]:
																								lowerCamelCase__								=	key_components[-2] + """_v"""
																		if name is not None:
																								lowerCamelCase__								=	key_components[:-3] + [name]
																								lowerCamelCase__								=	""".""".join(__lowerCAmelCase    )
																								lowerCamelCase__								=	key

												if flax_key in special_pt_names:
																		lowerCamelCase__								=	special_pt_names[flax_key]

												if flax_key in pt_model_dict:
																		if flax_tensor.shape != pt_model_dict[flax_key].shape:
																								raise ValueError(
																								    F'''Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '''
																								    F'''to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'''    )
																		else:
																								# add weight to pytorch dict
																								lowerCamelCase__								=	np.asarray(__lowerCAmelCase    ) if not isinstance(__lowerCAmelCase ,     np.ndarray    ) else flax_tensor
																								lowerCamelCase__								=	torch.from_numpy(__lowerCAmelCase    )
																								# remove from missing keys
																								missing_keys.remove(__lowerCAmelCase    )
												else:
																		# weight is not expected by PyTorch model
																		unexpected_keys.append(__lowerCAmelCase    )

						pt_model.load_state_dict(__lowerCAmelCase    )

						# re-transform missing_keys to list
						lowerCamelCase__								=	list(__lowerCAmelCase    )

						if len(__lowerCAmelCase    ) > 0:
												logger.warning(
												    """Some weights of the Flax model were not used when initializing the PyTorch model"""
												    F''' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'''
												    F''' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'''
												    """ (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"""
												    F''' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'''
												    """ to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"""
												    """ FlaxBertForSequenceClassification model)."""    )
						else:
												logger.warning(F'''All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n'''    )
						if len(__lowerCAmelCase    ) > 0:
												logger.warning(
												    F'''Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'''
												    F''' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'''
												    """ use it for predictions and inference."""    )
						else:
												logger.warning(
												    F'''All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n'''
												    """If your task is similar to the task the model of the checkpoint was trained on, """
												    F'''you can already use {pt_model.__class__.__name__} for predictions without further training.'''    )

						return pt_model



'''simple docstring'''




import collections
import inspect
import unittest
from typing import Dict, List, Tuple

from transformers import MaskFormerSwinConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, torch_device
from transformers.utils import is_torch_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 torch import nn

      from transformers import MaskFormerSwinBackbone
      from transformers.models.maskformer import MaskFormerSwinModel





class 					A							:





     def __init__(				self		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE=13		,     SCREAMING_SNAKE_CASE=32		,     SCREAMING_SNAKE_CASE=2		,     SCREAMING_SNAKE_CASE=3		,     SCREAMING_SNAKE_CASE=16		,     SCREAMING_SNAKE_CASE=[1, 2, 1]		,     SCREAMING_SNAKE_CASE=[2, 2, 4]		,     SCREAMING_SNAKE_CASE=2		,     SCREAMING_SNAKE_CASE=2.0		,     SCREAMING_SNAKE_CASE=True		,     SCREAMING_SNAKE_CASE=0.0		,     SCREAMING_SNAKE_CASE=0.0		,     SCREAMING_SNAKE_CASE=0.1		,     SCREAMING_SNAKE_CASE="gelu"		,     SCREAMING_SNAKE_CASE=False		,     SCREAMING_SNAKE_CASE=True		,     SCREAMING_SNAKE_CASE=0.02		,     SCREAMING_SNAKE_CASE=1e-5		,     SCREAMING_SNAKE_CASE=True		,     SCREAMING_SNAKE_CASE=None		,     SCREAMING_SNAKE_CASE=True		,     SCREAMING_SNAKE_CASE=10		,     SCREAMING_SNAKE_CASE=8		,     SCREAMING_SNAKE_CASE=["stage1", "stage2", "stage3"]		,     SCREAMING_SNAKE_CASE=[1, 2, 3]		,     )					->							Optional[int]:






         """simple docstring"""






         A			:  Any   =		parent
         A			:  List[Any]   =		batch_size
         A			:  Any   =		image_size
         A			:  Optional[Any]   =		patch_size
         A			:  Dict   =		num_channels
         A			:  Union[str, Any]   =		embed_dim
         A			:  Optional[int]   =		depths
         A			:  int   =		num_heads
         A			:  Tuple   =		window_size
         A			:  str   =		mlp_ratio
         A			:  Optional[Any]   =		qkv_bias
         A			:  Optional[int]   =		hidden_dropout_prob
         A			:  Any   =		attention_probs_dropout_prob
         A			:  List[Any]   =		drop_path_rate
         A			:  Tuple   =		hidden_act
         A			:  Any   =		use_absolute_embeddings
         A			:  Optional[int]   =		patch_norm
         A			:  List[str]   =		layer_norm_eps
         A			:  int   =		initializer_range
         A			:  Tuple   =		is_training
         A			:  int   =		scope
         A			:  List[Any]   =		use_labels
         A			:  List[str]   =		type_sequence_label_size
         A			:  int   =		encoder_stride
         A			:  List[Any]   =		out_features
         A			:  str   =		out_indices





     def 							__lowerCAmelCase				(				self      )					->							str:






         """simple docstring"""






         A			:  Optional[int]   =		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

         A			:  Union[str, Any]   =		None
         if self.use_labels:
             A			:  int   =		ids_tensor([self.batch_size]		,     self.type_sequence_label_size      )

         A			:  str   =		self.get_config()

         return config, pixel_values, labels





     def 							__lowerCAmelCase				(				self      )					->							List[Any]:






         """simple docstring"""






         return MaskFormerSwinConfig(
             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		,     window_size=self.window_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		,     use_absolute_embeddings=self.use_absolute_embeddings		,     path_norm=self.patch_norm		,     layer_norm_eps=self.layer_norm_eps		,     initializer_range=self.initializer_range		,     encoder_stride=self.encoder_stride		,     out_features=self.out_features		,     out_indices=self.out_indices		,     )





     def 							__lowerCAmelCase				(				self		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )					->							Dict:






         """simple docstring"""






         A			:  List[Any]   =		MaskFormerSwinModel(config=SCREAMING_SNAKE_CASE      )
         model.to(SCREAMING_SNAKE_CASE      )
         model.eval()
         A			:  Optional[int]   =		model(SCREAMING_SNAKE_CASE      )

         A			:  Tuple   =		((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths      ) - 1))
         A			:  Optional[int]   =		int(config.embed_dim * 2 ** (len(config.depths      ) - 1)      )

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





     def 							__lowerCAmelCase				(				self		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )					->							Any:






         """simple docstring"""






         A			:  Tuple   =		MaskFormerSwinBackbone(config=SCREAMING_SNAKE_CASE      )
         model.to(SCREAMING_SNAKE_CASE      )
         model.eval()
         A			:  Union[str, Any]   =		model(SCREAMING_SNAKE_CASE      )

         # verify feature maps
         self.parent.assertEqual(len(result.feature_maps      )		,     len(config.out_features      )      )
         self.parent.assertListEqual(list(result.feature_maps[0].shape      )		,     [13, 16, 16, 16]      )

         # verify channels
         self.parent.assertEqual(len(model.channels      )		,     len(config.out_features      )      )
         self.parent.assertListEqual(model.channels		,     [16, 32, 64]      )

         # verify ValueError
         with self.parent.assertRaises(SCREAMING_SNAKE_CASE      ):
             A			:  str   =		['''stem''']
             A			:  Dict   =		MaskFormerSwinBackbone(config=SCREAMING_SNAKE_CASE      )





     def 							__lowerCAmelCase				(				self      )					->							List[str]:






         """simple docstring"""






         A			:  str   =		self.prepare_config_and_inputs()
         A,    A,    A			:  int   =		config_and_inputs
         A			:  List[str]   =		{'''pixel_values''': pixel_values}
         return config, inputs_dict





@require_torch
class 					A							(				__snake_case       ,     __snake_case       ,     unittest.TestCase		):
     __magic_name__          =						(
         (
             MaskFormerSwinModel,
             MaskFormerSwinBackbone,
         )
         if is_torch_available()
         else ()
     )
     __magic_name__          =						{'''feature-extraction''': MaskFormerSwinModel} if is_torch_available() else {}
     __magic_name__          =						False
     __magic_name__          =						False
     __magic_name__          =						False
     __magic_name__          =						False
     __magic_name__          =						False





     def 							__lowerCAmelCase				(				self      )					->							Optional[int]:






         """simple docstring"""






         A			:  Tuple   =		MaskFormerSwinModelTester(self      )
         A			:  int   =		ConfigTester(self		,     config_class=SCREAMING_SNAKE_CASE		,     embed_dim=37      )





     @require_torch_multi_gpu
     @unittest.skip(
         reason=(
             '''`MaskFormerSwinModel` outputs `hidden_states_spatial_dimensions` which doesn\'t work well with'''
             ''' `nn.DataParallel`'''
         )      )
     def 							__lowerCAmelCase				(				self      )					->							Any:






         """simple docstring"""






         pass





     def 							__lowerCAmelCase				(				self      )					->							List[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 							__lowerCAmelCase				(				self      )					->							Dict:






         """simple docstring"""






         return





     def 							__lowerCAmelCase				(				self      )					->							Dict:






         """simple docstring"""






         A			:  Optional[Any]   =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE      )





     def 							__lowerCAmelCase				(				self      )					->							int:






         """simple docstring"""






         A			:  Optional[Any]   =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_backbone(*SCREAMING_SNAKE_CASE      )





     @unittest.skip('''Swin does not use inputs_embeds'''      )
     def 							__lowerCAmelCase				(				self      )					->							Optional[Any]:






         """simple docstring"""






         pass





     @unittest.skip('''Swin does not support feedforward chunking'''      )
     def 							__lowerCAmelCase				(				self      )					->							Optional[Any]:






         """simple docstring"""






         pass





     def 							__lowerCAmelCase				(				self      )					->							Any:






         """simple docstring"""






         A,    A			:  List[str]   =		self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             A			:  Optional[int]   =		model_class(SCREAMING_SNAKE_CASE      )
             self.assertIsInstance(model.get_input_embeddings()		,     (nn.Module)      )
             A			:  int   =		model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(SCREAMING_SNAKE_CASE		,     nn.Linear      )      )





     def 							__lowerCAmelCase				(				self      )					->							Dict:






         """simple docstring"""






         A,    A			:  List[str]   =		self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             A			:  List[Any]   =		model_class(SCREAMING_SNAKE_CASE      )
             A			:  List[str]   =		inspect.signature(model.forward      )
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             A			:  Optional[Any]   =		[*signature.parameters.keys()]

             A			:  Union[str, Any]   =		['''pixel_values''']
             self.assertListEqual(arg_names[:1]		,     SCREAMING_SNAKE_CASE      )





     @unittest.skip(reason='''MaskFormerSwin is only used as backbone and doesn\'t support output_attentions'''      )
     def 							__lowerCAmelCase				(				self      )					->							Any:






         """simple docstring"""






         pass





     @unittest.skip(reason='''MaskFormerSwin is only used as an internal backbone'''      )
     def 							__lowerCAmelCase				(				self      )					->							Tuple:






         """simple docstring"""






         pass





     def 							__lowerCAmelCase				(				self		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )					->							Optional[int]:






         """simple docstring"""






         A			:  Tuple   =		model_class(SCREAMING_SNAKE_CASE      )
         model.to(SCREAMING_SNAKE_CASE      )
         model.eval()

         with torch.no_grad():
             A			:  Any   =		model(**self._prepare_for_class(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )      )

         A			:  List[str]   =		outputs.hidden_states

         A			:  str   =		getattr(
             self.model_tester		,     '''expected_num_hidden_layers'''		,     len(self.model_tester.depths      ) + 1      )
         self.assertEqual(len(SCREAMING_SNAKE_CASE      )		,     SCREAMING_SNAKE_CASE      )

         # Swin has a different seq_length
         A			:  List[str]   =		(
             config.patch_size
             if isinstance(config.patch_size		,     collections.abc.Iterable      )
             else (config.patch_size, config.patch_size)
         )

         A			:  Optional[Any]   =		(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

         self.assertListEqual(
             list(hidden_states[0].shape[-2:]      )		,     [num_patches, self.model_tester.embed_dim]		,     )





     def 							__lowerCAmelCase				(				self      )					->							Any:






         """simple docstring"""






         A,    A			:  Union[str, Any]   =		self.model_tester.prepare_config_and_inputs_for_common()

         A			:  Union[str, Any]   =		(
             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:
             A			:  str   =		True
             self.check_hidden_states_output(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )

             # check that output_hidden_states also work using config
             del inputs_dict["output_hidden_states"]
             A			:  Optional[Any]   =		True

             self.check_hidden_states_output(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )





     def 							__lowerCAmelCase				(				self      )					->							Dict:






         """simple docstring"""






         A,    A			:  Any   =		self.model_tester.prepare_config_and_inputs_for_common()
         A			:  Union[str, Any]   =		3

         A			:  List[str]   =		(
             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)
         )
         A			:  Dict   =		(
             config.patch_size
             if isinstance(config.patch_size		,     collections.abc.Iterable      )
             else (config.patch_size, config.patch_size)
         )

         A			:  Any   =		image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
         A			:  Any   =		image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

         for model_class in self.all_model_classes:
             A			:  Optional[Any]   =		True
             self.check_hidden_states_output(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     (padded_height, padded_width)      )

             # check that output_hidden_states also work using config
             del inputs_dict["output_hidden_states"]
             A			:  Optional[Any]   =		True
             self.check_hidden_states_output(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     (padded_height, padded_width)      )





     @unittest.skip(reason='''MaskFormerSwin doesn\'t have pretrained checkpoints'''      )
     def 							__lowerCAmelCase				(				self      )					->							Dict:






         """simple docstring"""






         pass





     @unittest.skip(reason='''This will be fixed once MaskFormerSwin is replaced by native Swin'''      )
     def 							__lowerCAmelCase				(				self      )					->							int:






         """simple docstring"""






         pass





     @unittest.skip(reason='''This will be fixed once MaskFormerSwin is replaced by native Swin'''      )
     def 							__lowerCAmelCase				(				self      )					->							Optional[int]:






         """simple docstring"""






         pass





     def 							__lowerCAmelCase				(				self      )					->							List[Any]:






         """simple docstring"""






         A,    A			:  str   =		self.model_tester.prepare_config_and_inputs_for_common()

         def set_nan_tensor_to_zero(SCREAMING_SNAKE_CASE      ):
             A			:  List[str]   =		0
             return t

         def check_equivalence(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE={}      ):
             with torch.no_grad():
                 A			:  List[Any]   =		model(**SCREAMING_SNAKE_CASE		,     return_dict=SCREAMING_SNAKE_CASE		,     **SCREAMING_SNAKE_CASE      )
                 A			:  str   =		model(**SCREAMING_SNAKE_CASE		,     return_dict=SCREAMING_SNAKE_CASE		,     **SCREAMING_SNAKE_CASE      ).to_tuple()

                 def recursive_check(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      ):
                     if isinstance(SCREAMING_SNAKE_CASE		,     (List, Tuple)      ):
                         for tuple_iterable_value, dict_iterable_value in zip(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      ):
                             recursive_check(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )
                     elif isinstance(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      ):
                         for tuple_iterable_value, dict_iterable_value in zip(
                             tuple_object.values()		,     dict_object.values()      ):
                             recursive_check(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )
                     elif tuple_object is None:
                         return
                     else:
                         self.assertTrue(
                             torch.allclose(
                                 set_nan_tensor_to_zero(SCREAMING_SNAKE_CASE      )		,     set_nan_tensor_to_zero(SCREAMING_SNAKE_CASE      )		,     atol=1e-5      )		,     msg=(
                                 '''Tuple and dict output are not equal. Difference:'''
                                 F' {torch.max(torch.abs(tuple_object - dict_object      )      )}. Tuple has `nan`:'
                                 F' {torch.isnan(SCREAMING_SNAKE_CASE      ).any()} and `inf`: {torch.isinf(SCREAMING_SNAKE_CASE      )}. Dict has'
                                 F' `nan`: {torch.isnan(SCREAMING_SNAKE_CASE      ).any()} and `inf`: {torch.isinf(SCREAMING_SNAKE_CASE      )}.'
                             )		,     )

                 recursive_check(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )

         for model_class in self.all_model_classes:
             A			:  Tuple   =		model_class(SCREAMING_SNAKE_CASE      )
             model.to(SCREAMING_SNAKE_CASE      )
             model.eval()

             A			:  List[str]   =		self._prepare_for_class(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )
             A			:  int   =		self._prepare_for_class(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )
             check_equivalence(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )

             A			:  Union[str, Any]   =		self._prepare_for_class(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     return_labels=SCREAMING_SNAKE_CASE      )
             A			:  int   =		self._prepare_for_class(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     return_labels=SCREAMING_SNAKE_CASE      )
             check_equivalence(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )

             A			:  Tuple   =		self._prepare_for_class(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )
             A			:  Any   =		self._prepare_for_class(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )
             check_equivalence(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     {'''output_hidden_states''': True}      )

             A			:  Optional[int]   =		self._prepare_for_class(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     return_labels=SCREAMING_SNAKE_CASE      )
             A			:  Dict   =		self._prepare_for_class(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     return_labels=SCREAMING_SNAKE_CASE      )
             check_equivalence(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     {'''output_hidden_states''': True}      )





@require_torch
class 					A							(				unittest.TestCase       ,     __snake_case		):
     __magic_name__          =						(MaskFormerSwinBackbone,) if is_torch_available() else ()
     __magic_name__          =						MaskFormerSwinConfig





     def 							__lowerCAmelCase				(				self      )					->							Any:






         """simple docstring"""






         A			:  Dict   =		MaskFormerSwinModelTester(self      )





     def 							__lowerCAmelCase				(				self      )					->							Dict:






         """simple docstring"""






         A,    A			:  int   =		self.model_tester.prepare_config_and_inputs_for_common()
         A			:  str   =		inputs_dict['''pixel_values'''].shape[0]

         for backbone_class in self.all_model_classes:
             A			:  Any   =		backbone_class(SCREAMING_SNAKE_CASE      )
             backbone.to(SCREAMING_SNAKE_CASE      )
             backbone.eval()

             A			:  Dict   =		backbone(**SCREAMING_SNAKE_CASE      )

             # Test default outputs and verify feature maps
             self.assertIsInstance(outputs.feature_maps		,     SCREAMING_SNAKE_CASE      )
             self.assertTrue(len(outputs.feature_maps      ) == len(backbone.channels      )      )
             for feature_map, n_channels in zip(outputs.feature_maps		,     backbone.channels      ):
                 self.assertTrue(feature_map.shape[:2]		,     (batch_size, n_channels)      )
             self.assertIsNone(outputs.hidden_states      )
             self.assertIsNone(outputs.attentions      )

             # Test output_hidden_states=True
             A			:  Optional[int]   =		backbone(**SCREAMING_SNAKE_CASE		,     output_hidden_states=SCREAMING_SNAKE_CASE      )
             self.assertIsNotNone(outputs.hidden_states      )
             self.assertTrue(len(outputs.hidden_states      )		,     len(backbone.stage_names      )      )
             # We skip the stem layer
             for hidden_states, n_channels in zip(outputs.hidden_states[1:]		,     backbone.channels      ):
                 for hidden_state in hidden_states:
                     # Hidden states are in the format (batch_size, (height * width), n_channels)
                     A,    A,    A			:  List[Any]   =		hidden_state.shape
                     self.assertTrue((h_batch_size, h_n_channels)		,     (batch_size, n_channels)      )

            # Test output_attentions=True
             if self.has_attentions:
                 A			:  Dict   =		backbone(**SCREAMING_SNAKE_CASE		,     output_attentions=SCREAMING_SNAKE_CASE      )
                 self.assertIsNotNone(outputs.attentions      )




