
from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_outputs import (
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import logging
from .configuration_regnet import RegNetConfig


lowercase__		:Union[str, Any]   =    logging.get_logger(__name__)

# General docstring
lowercase__		:int   =    "RegNetConfig"

# Base docstring
lowercase__		:Optional[int]   =    "facebook/regnet-y-040"
lowercase__		:List[str]   =    [1, 1088, 7, 7]

# Image classification docstring
lowercase__		:Dict   =    "facebook/regnet-y-040"
lowercase__		:Union[str, Any]   =    "tabby, tabby cat"

lowercase__		:List[str]   =    [
    "facebook/regnet-y-040",
    # See all regnet models at https://huggingface.co/models?filter=regnet
]







class 				lowercase		(					nn.Module		):




							def __init__(				self   ,A__   ,A__   ,A__ = 3   ,A__ = 1   ,A__ = 1   ,A__ = "relu"   ,):
												super().__init__()
												lowercase							=						nn.Convad(
												    __snake_case   ,__snake_case   ,kernel_size=__snake_case   ,stride=__snake_case   ,padding=kernel_size // 2   ,groups=__snake_case   ,bias=__snake_case   ,)
												lowercase							=						nn.BatchNormad(__snake_case)
												lowercase							=						ACTaFN[activation] if activation is not None else nn.Identity()





							def  A__      (				self   ,A__):
												lowercase							=						self.convolution(__snake_case)
												lowercase							=						self.normalization(__snake_case)
												lowercase							=						self.activation(__snake_case)
												return hidden_state







class 				lowercase		(					nn.Module		):




							def __init__(				self   ,A__):
												super().__init__()
												lowercase							=						RegNetConvLayer(
												    config.num_channels   ,config.embedding_size   ,kernel_size=3   ,stride=2   ,activation=config.hidden_act)
												lowercase							=						config.num_channels





							def  A__      (				self   ,A__):
												lowercase							=						pixel_values.shape[1]
												if num_channels != self.num_channels:
																	raise ValueError(
																	    '''Make sure that the channel dimension of the pixel values match with the one set in the configuration.''')
												lowercase							=						self.embedder(__snake_case)
												return hidden_state







class 				lowercase		(					nn.Module		):




							def __init__(				self   ,A__   ,A__   ,A__ = 2):
												super().__init__()
												lowercase							=						nn.Convad(__snake_case   ,__snake_case   ,kernel_size=1   ,stride=__snake_case   ,bias=__snake_case)
												lowercase							=						nn.BatchNormad(__snake_case)





							def  A__      (				self   ,A__):
												lowercase							=						self.convolution(__snake_case)
												lowercase							=						self.normalization(__snake_case)
												return hidden_state







class 				lowercase		(					nn.Module		):




							def __init__(				self   ,A__   ,A__):
												super().__init__()

												lowercase							=						nn.AdaptiveAvgPoolad((1, 1))
												lowercase							=						nn.Sequential(
												    nn.Convad(__snake_case   ,__snake_case   ,kernel_size=1)   ,nn.ReLU()   ,nn.Convad(__snake_case   ,__snake_case   ,kernel_size=1)   ,nn.Sigmoid()   ,)





							def  A__      (				self   ,A__):
												# b c h w -> b c 1 1
												lowercase							=						self.pooler(__snake_case)
												lowercase							=						self.attention(__snake_case)
												lowercase							=						hidden_state * attention
												return hidden_state







class 				lowercase		(					nn.Module		):




							def __init__(				self   ,A__   ,A__   ,A__   ,A__ = 1):
												super().__init__()
												lowercase							=						in_channels != out_channels or stride != 1
												lowercase							=						max(1   ,out_channels // config.groups_width)
												lowercase							=						(
												    RegNetShortCut(__snake_case   ,__snake_case   ,stride=__snake_case) if should_apply_shortcut else nn.Identity()
												)
												lowercase							=						nn.Sequential(
												    RegNetConvLayer(__snake_case   ,__snake_case   ,kernel_size=1   ,activation=config.hidden_act)   ,RegNetConvLayer(__snake_case   ,__snake_case   ,stride=__snake_case   ,groups=__snake_case   ,activation=config.hidden_act)   ,RegNetConvLayer(__snake_case   ,__snake_case   ,kernel_size=1   ,activation=__snake_case)   ,)
												lowercase							=						ACTaFN[config.hidden_act]





							def  A__      (				self   ,A__):
												lowercase							=						hidden_state
												lowercase							=						self.layer(__snake_case)
												lowercase							=						self.shortcut(__snake_case)
												hidden_state += residual
												lowercase							=						self.activation(__snake_case)
												return hidden_state







class 				lowercase		(					nn.Module		):




							def __init__(				self   ,A__   ,A__   ,A__   ,A__ = 1):
												super().__init__()
												lowercase							=						in_channels != out_channels or stride != 1
												lowercase							=						max(1   ,out_channels // config.groups_width)
												lowercase							=						(
												    RegNetShortCut(__snake_case   ,__snake_case   ,stride=__snake_case) if should_apply_shortcut else nn.Identity()
												)
												lowercase							=						nn.Sequential(
												    RegNetConvLayer(__snake_case   ,__snake_case   ,kernel_size=1   ,activation=config.hidden_act)   ,RegNetConvLayer(__snake_case   ,__snake_case   ,stride=__snake_case   ,groups=__snake_case   ,activation=config.hidden_act)   ,RegNetSELayer(__snake_case   ,reduced_channels=int(round(in_channels / 4)))   ,RegNetConvLayer(__snake_case   ,__snake_case   ,kernel_size=1   ,activation=__snake_case)   ,)
												lowercase							=						ACTaFN[config.hidden_act]





							def  A__      (				self   ,A__):
												lowercase							=						hidden_state
												lowercase							=						self.layer(__snake_case)
												lowercase							=						self.shortcut(__snake_case)
												hidden_state += residual
												lowercase							=						self.activation(__snake_case)
												return hidden_state







class 				lowercase		(					nn.Module		):




							def __init__(				self   ,A__   ,A__   ,A__   ,A__ = 2   ,A__ = 2   ,):
												super().__init__()

												lowercase							=						RegNetXLayer if config.layer_type == '''x''' else RegNetYLayer

												lowercase							=						nn.Sequential(
												    # downsampling is done in the first layer with stride of 2
												    layer(
												        __snake_case   ,__snake_case   ,__snake_case   ,stride=__snake_case   ,)   ,*[layer(__snake_case   ,__snake_case   ,__snake_case) for _ in range(depth - 1)]   ,)





							def  A__      (				self   ,A__):
												lowercase							=						self.layers(__snake_case)
												return hidden_state







class 				lowercase		(					nn.Module		):




							def __init__(				self   ,A__):
												super().__init__()
												lowercase							=						nn.ModuleList([])
												# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
												self.stages.append(
												    RegNetStage(
												        __snake_case   ,config.embedding_size   ,config.hidden_sizes[0]   ,stride=2 if config.downsample_in_first_stage else 1   ,depth=config.depths[0]   ,))
												lowercase							=						zip(config.hidden_sizes   ,config.hidden_sizes[1:])
												for (in_channels, out_channels), depth in zip(__snake_case   ,config.depths[1:]):
																	self.stages.append(RegNetStage(__snake_case   ,__snake_case   ,__snake_case   ,depth=__snake_case))





							def  A__      (				self   ,A__   ,A__ = False   ,A__ = True):
												lowercase							=						() if output_hidden_states else None

												for stage_module in self.stages:
																	if output_hidden_states:
																						lowercase							=						hidden_states + (hidden_state,)

																	lowercase							=						stage_module(__snake_case)

												if output_hidden_states:
																	lowercase							=						hidden_states + (hidden_state,)

												if not return_dict:
																	return tuple(v for v in [hidden_state, hidden_states] if v is not None)

												return BaseModelOutputWithNoAttention(last_hidden_state=__snake_case   ,hidden_states=__snake_case)







class 				lowercase		(					A__		):
							lowercase_       :       Union[str, Any]       					=RegNetConfig
							lowercase_       :       Optional[int]       					='''regnet'''
							lowercase_       :       Optional[int]       					='''pixel_values'''
							lowercase_       :       int       					=True




							def  A__      (				self   ,A__):
												if isinstance(__snake_case   ,nn.Convad):
																	nn.init.kaiming_normal_(module.weight   ,mode='''fan_out'''   ,nonlinearity='''relu''')
												elif isinstance(__snake_case   ,(nn.BatchNormad, nn.GroupNorm)):
																	nn.init.constant_(module.weight   ,1)
																	nn.init.constant_(module.bias   ,0)





							def  A__      (				self   ,A__   ,A__=False):
												if isinstance(__snake_case   ,__snake_case):
																	lowercase							=						value


lowercase__		:List[str]   =    r"\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"

lowercase__		:Any   =    r"\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConvNextImageProcessor.__call__`] for details.\n\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.\n"







@add_start_docstrings(
    '''The bare RegNet model outputting raw features without any specific head on top.'''       ,  A__       ,  )
# Copied from transformers.models.resnet.modeling_resnet.ResNetModel with RESNET->REGNET,ResNet->RegNet
class 				lowercase		(					A__		):




							def __init__(				self   ,A__):
												super().__init__(__snake_case)
												lowercase							=						config
												lowercase							=						RegNetEmbeddings(__snake_case)
												lowercase							=						RegNetEncoder(__snake_case)
												lowercase							=						nn.AdaptiveAvgPoolad((1, 1))
												# Initialize weights and apply final processing
												self.post_init()





							@add_start_docstrings_to_model_forward(__snake_case)
							@add_code_sample_docstrings(
							    checkpoint=_CHECKPOINT_FOR_DOC   ,output_type=__snake_case   ,config_class=_CONFIG_FOR_DOC   ,modality='''vision'''   ,expected_output=_EXPECTED_OUTPUT_SHAPE   ,)
							def  A__      (				self   ,A__   ,A__ = None   ,A__ = None):
												lowercase							=						(
												    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
												)
												lowercase							=						return_dict if return_dict is not None else self.config.use_return_dict

												lowercase							=						self.embedder(__snake_case)

												lowercase							=						self.encoder(
												    __snake_case   ,output_hidden_states=__snake_case   ,return_dict=__snake_case)

												lowercase							=						encoder_outputs[0]

												lowercase							=						self.pooler(__snake_case)

												if not return_dict:
																	return (last_hidden_state, pooled_output) + encoder_outputs[1:]

												return BaseModelOutputWithPoolingAndNoAttention(
												    last_hidden_state=__snake_case   ,pooler_output=__snake_case   ,hidden_states=encoder_outputs.hidden_states   ,)







@add_start_docstrings(
    '''
    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    '''       ,  A__       ,  )
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class 				lowercase		(					A__		):




							def __init__(				self   ,A__):
												super().__init__(__snake_case)
												lowercase							=						config.num_labels
												lowercase							=						RegNetModel(__snake_case)
												# classification head
												lowercase							=						nn.Sequential(
												    nn.Flatten()   ,nn.Linear(config.hidden_sizes[-1]   ,config.num_labels) if config.num_labels > 0 else nn.Identity()   ,)
												# initialize weights and apply final processing
												self.post_init()





							@add_start_docstrings_to_model_forward(__snake_case)
							@add_code_sample_docstrings(
							    checkpoint=_IMAGE_CLASS_CHECKPOINT   ,output_type=__snake_case   ,config_class=_CONFIG_FOR_DOC   ,expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT   ,)
							def  A__      (				self   ,A__ = None   ,A__ = None   ,A__ = None   ,A__ = None   ,):
												lowercase							=						return_dict if return_dict is not None else self.config.use_return_dict

												lowercase							=						self.regnet(__snake_case   ,output_hidden_states=__snake_case   ,return_dict=__snake_case)

												lowercase							=						outputs.pooler_output if return_dict else outputs[1]

												lowercase							=						self.classifier(__snake_case)

												lowercase							=						None

												if labels is not None:
																	if self.config.problem_type is None:
																						if self.num_labels == 1:
																											lowercase							=						'''regression'''
																						elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																											lowercase							=						'''single_label_classification'''
																						else:
																											lowercase							=						'''multi_label_classification'''
																	if self.config.problem_type == "regression":
																						lowercase							=						MSELoss()
																						if self.num_labels == 1:
																											lowercase							=						loss_fct(logits.squeeze()   ,labels.squeeze())
																						else:
																											lowercase							=						loss_fct(__snake_case   ,__snake_case)
																	elif self.config.problem_type == "single_label_classification":
																						lowercase							=						CrossEntropyLoss()
																						lowercase							=						loss_fct(logits.view(-1   ,self.num_labels)   ,labels.view(-1))
																	elif self.config.problem_type == "multi_label_classification":
																						lowercase							=						BCEWithLogitsLoss()
																						lowercase							=						loss_fct(__snake_case   ,__snake_case)

												if not return_dict:
																	lowercase							=						(logits,) + outputs[2:]
																	return (loss,) + output if loss is not None else output

												return ImageClassifierOutputWithNoAttention(loss=__snake_case   ,logits=__snake_case   ,hidden_states=outputs.hidden_states)









'''simple docstring'''






import unittest

from transformers import MraConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_torch_available():
						import torch

						from transformers import (
						    MraForMaskedLM,
						    MraForMultipleChoice,
						    MraForQuestionAnswering,
						    MraForSequenceClassification,
						    MraForTokenClassification,
						    MraModel,
						)
						from transformers.models.mra.modeling_mra import MRA_PRETRAINED_MODEL_ARCHIVE_LIST

class 				SCREAMING_SNAKE_CASE:







	"""simple docstring"""






	def __init__(						self			:     Optional[int]		,			__snake_case			:     str		,			__snake_case			:     Union[str, Any]=2		,			__snake_case			:     Optional[int]=8		,			__snake_case			:     Any=True		,			__snake_case			:     Union[str, Any]=True		,			__snake_case			:     Dict=True		,			__snake_case			:     int=True		,			__snake_case			:     List[Any]=99		,			__snake_case			:     str=16		,			__snake_case			:     Tuple=5		,			__snake_case			:     Tuple=2		,			__snake_case			:     str=36		,			__snake_case			:     Dict="gelu"		,			__snake_case			:     str=0.0		,			__snake_case			:     Optional[int]=0.0		,			__snake_case			:     Optional[int]=512		,			__snake_case			:     Optional[Any]=16		,			__snake_case			:     int=2		,			__snake_case			:     int=0.02		,			__snake_case			:     str=3		,			__snake_case			:     Dict=4		,			__snake_case			:     str=None		,			)			->		Optional[int]:
								UpperCAmelCase :       Optional[int] 							=		parent
								UpperCAmelCase :       Tuple 							=		batch_size
								UpperCAmelCase :       List[str] 							=		seq_length
								UpperCAmelCase :       List[Any] 							=		is_training
								UpperCAmelCase :       int 							=		use_input_mask
								UpperCAmelCase :       Any 							=		use_token_type_ids
								UpperCAmelCase :       str 							=		use_labels
								UpperCAmelCase :       Union[str, Any] 							=		vocab_size
								UpperCAmelCase :       List[str] 							=		hidden_size
								UpperCAmelCase :       Optional[Any] 							=		num_hidden_layers
								UpperCAmelCase :       Union[str, Any] 							=		num_attention_heads
								UpperCAmelCase :       Optional[Any] 							=		intermediate_size
								UpperCAmelCase :       Union[str, Any] 							=		hidden_act
								UpperCAmelCase :       int 							=		hidden_dropout_prob
								UpperCAmelCase :       Optional[int] 							=		attention_probs_dropout_prob
								UpperCAmelCase :       Union[str, Any] 							=		max_position_embeddings
								UpperCAmelCase :       str 							=		type_vocab_size
								UpperCAmelCase :       List[str] 							=		type_sequence_label_size
								UpperCAmelCase :       Tuple 							=		initializer_range
								UpperCAmelCase :       Optional[Any] 							=		num_labels
								UpperCAmelCase :       Optional[int] 							=		num_choices
								UpperCAmelCase :       Any 							=		scope







	def      A	(						self			:     Tuple		)			->		Union[str, Any]:
								UpperCAmelCase :       Optional[int] 							=		ids_tensor([self.batch_size, self.seq_length]		,			self.vocab_size		)

								UpperCAmelCase :       Optional[int] 							=		None
								if self.use_input_mask:
															UpperCAmelCase :       Dict 							=		random_attention_mask([self.batch_size, self.seq_length]		)

								UpperCAmelCase :       Dict 							=		None
								if self.use_token_type_ids:
															UpperCAmelCase :       Union[str, Any] 							=		ids_tensor([self.batch_size, self.seq_length]		,			self.type_vocab_size		)

								UpperCAmelCase :       str 							=		None
								UpperCAmelCase :       Tuple 							=		None
								UpperCAmelCase :       int 							=		None
								if self.use_labels:
															UpperCAmelCase :       str 							=		ids_tensor([self.batch_size]		,			self.type_sequence_label_size		)
															UpperCAmelCase :       Dict 							=		ids_tensor([self.batch_size, self.seq_length]		,			self.num_labels		)
															UpperCAmelCase :       str 							=		ids_tensor([self.batch_size]		,			self.num_choices		)

								UpperCAmelCase :       List[str] 							=		self.get_config()

								return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







	def      A	(						self			:     int		)			->		Tuple:
								return MraConfig(
								    vocab_size=self.vocab_size		,			hidden_size=self.hidden_size		,			num_hidden_layers=self.num_hidden_layers		,			num_attention_heads=self.num_attention_heads		,			intermediate_size=self.intermediate_size		,			hidden_act=self.hidden_act		,			hidden_dropout_prob=self.hidden_dropout_prob		,			attention_probs_dropout_prob=self.attention_probs_dropout_prob		,			max_position_embeddings=self.max_position_embeddings		,			type_vocab_size=self.type_vocab_size		,			is_decoder=__snake_case		,			initializer_range=self.initializer_range		,			)







	def      A	(						self			:     Optional[Any]		)			->		Any:
								UpperCAmelCase :       Optional[Any] 							=		self.get_config()
								UpperCAmelCase :       int 							=		300
								return config







	def      A	(						self			:     Optional[Any]		)			->		Any:
								(
								    (
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		
								) :       Tuple 							=		self.prepare_config_and_inputs()

								UpperCAmelCase :       Dict 							=		True
								UpperCAmelCase :       Tuple 							=		floats_tensor([self.batch_size, self.seq_length, self.hidden_size]		)
								UpperCAmelCase :       str 							=		ids_tensor([self.batch_size, self.seq_length]		,			vocab_size=2		)

								return (
								    config,
								    input_ids,
								    token_type_ids,
								    input_mask,
								    sequence_labels,
								    token_labels,
								    choice_labels,
								    encoder_hidden_states,
								    encoder_attention_mask,
								)







	def      A	(						self			:     Union[str, Any]		,			__snake_case			:     Union[str, Any]		,			__snake_case			:     int		,			__snake_case			:     Optional[int]		,			__snake_case			:     int		,			__snake_case			:     Dict		,			__snake_case			:     Tuple		,			__snake_case			:     Optional[Any]		)			->		List[str]:
								UpperCAmelCase :       int 							=		MraModel(config=__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       Tuple 							=		model(__snake_case		,			attention_mask=__snake_case		,			token_type_ids=__snake_case		)
								UpperCAmelCase :       Optional[int] 							=		model(__snake_case		,			token_type_ids=__snake_case		)
								UpperCAmelCase :       Dict 							=		model(__snake_case		)
								self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, self.seq_length, self.hidden_size)		)







	def      A	(						self			:     Tuple		,			__snake_case			:     Optional[Any]		,			__snake_case			:     List[str]		,			__snake_case			:     List[str]		,			__snake_case			:     int		,			__snake_case			:     Union[str, Any]		,			__snake_case			:     Optional[Any]		,			__snake_case			:     Any		,			__snake_case			:     List[Any]		,			__snake_case			:     Optional[Any]		,			)			->		Tuple:
								UpperCAmelCase :       str 							=		True
								UpperCAmelCase :       Tuple 							=		MraModel(__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       Optional[int] 							=		model(
								    __snake_case		,			attention_mask=__snake_case		,			token_type_ids=__snake_case		,			encoder_hidden_states=__snake_case		,			encoder_attention_mask=__snake_case		,			)
								UpperCAmelCase :       Optional[Any] 							=		model(
								    __snake_case		,			attention_mask=__snake_case		,			token_type_ids=__snake_case		,			encoder_hidden_states=__snake_case		,			)
								UpperCAmelCase :       str 							=		model(__snake_case		,			attention_mask=__snake_case		,			token_type_ids=__snake_case		)
								self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, self.seq_length, self.hidden_size)		)







	def      A	(						self			:     Tuple		,			__snake_case			:     str		,			__snake_case			:     Optional[Any]		,			__snake_case			:     List[str]		,			__snake_case			:     Tuple		,			__snake_case			:     Optional[Any]		,			__snake_case			:     List[str]		,			__snake_case			:     int		)			->		Any:
								UpperCAmelCase :       Dict 							=		MraForMaskedLM(config=__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       Optional[int] 							=		model(__snake_case		,			attention_mask=__snake_case		,			token_type_ids=__snake_case		,			labels=__snake_case		)
								self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.seq_length, self.vocab_size)		)







	def      A	(						self			:     Tuple		,			__snake_case			:     Tuple		,			__snake_case			:     Dict		,			__snake_case			:     Dict		,			__snake_case			:     Any		,			__snake_case			:     int		,			__snake_case			:     Optional[Any]		,			__snake_case			:     Tuple		)			->		Optional[int]:
								UpperCAmelCase :       List[str] 							=		MraForQuestionAnswering(config=__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       List[Any] 							=		model(
								    __snake_case		,			attention_mask=__snake_case		,			token_type_ids=__snake_case		,			start_positions=__snake_case		,			end_positions=__snake_case		,			)
								self.parent.assertEqual(result.start_logits.shape		,			(self.batch_size, self.seq_length)		)
								self.parent.assertEqual(result.end_logits.shape		,			(self.batch_size, self.seq_length)		)







	def      A	(						self			:     str		,			__snake_case			:     Optional[int]		,			__snake_case			:     List[str]		,			__snake_case			:     str		,			__snake_case			:     int		,			__snake_case			:     Optional[Any]		,			__snake_case			:     Union[str, Any]		,			__snake_case			:     List[Any]		)			->		int:
								UpperCAmelCase :       int 							=		self.num_labels
								UpperCAmelCase :       Union[str, Any] 							=		MraForSequenceClassification(__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       List[str] 							=		model(__snake_case		,			attention_mask=__snake_case		,			token_type_ids=__snake_case		,			labels=__snake_case		)
								self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.num_labels)		)







	def      A	(						self			:     str		,			__snake_case			:     Dict		,			__snake_case			:     Optional[Any]		,			__snake_case			:     Dict		,			__snake_case			:     Union[str, Any]		,			__snake_case			:     Union[str, Any]		,			__snake_case			:     Union[str, Any]		,			__snake_case			:     Dict		)			->		int:
								UpperCAmelCase :       Tuple 							=		self.num_labels
								UpperCAmelCase :       List[str] 							=		MraForTokenClassification(config=__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       str 							=		model(__snake_case		,			attention_mask=__snake_case		,			token_type_ids=__snake_case		,			labels=__snake_case		)
								self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.seq_length, self.num_labels)		)







	def      A	(						self			:     str		,			__snake_case			:     int		,			__snake_case			:     Any		,			__snake_case			:     Tuple		,			__snake_case			:     Optional[int]		,			__snake_case			:     List[str]		,			__snake_case			:     str		,			__snake_case			:     Union[str, Any]		)			->		Optional[Any]:
								UpperCAmelCase :       Tuple 							=		self.num_choices
								UpperCAmelCase :       int 							=		MraForMultipleChoice(config=__snake_case		)
								model.to(__snake_case		)
								model.eval()
								UpperCAmelCase :       str 							=		input_ids.unsqueeze(1		).expand(-1		,			self.num_choices		,			-1		).contiguous()
								UpperCAmelCase :       List[Any] 							=		token_type_ids.unsqueeze(1		).expand(-1		,			self.num_choices		,			-1		).contiguous()
								UpperCAmelCase :       List[str] 							=		input_mask.unsqueeze(1		).expand(-1		,			self.num_choices		,			-1		).contiguous()
								UpperCAmelCase :       List[str] 							=		model(
								    __snake_case		,			attention_mask=__snake_case		,			token_type_ids=__snake_case		,			labels=__snake_case		,			)
								self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.num_choices)		)







	def      A	(						self			:     str		)			->		Dict:
								UpperCAmelCase :       Any 							=		self.prepare_config_and_inputs()
								(
								    (
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		(
								    UpperCAmelCase
								)							,		
								) :       List[str] 							=		config_and_inputs
								UpperCAmelCase :       Any 							=		{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
								return config, inputs_dict

@require_torch
class 				SCREAMING_SNAKE_CASE(							A__							,		unittest.TestCase			):







	"""simple docstring"""






	lowerCamelCase__			 =       (
	    (
	        MraModel,
	        MraForMaskedLM,
	        MraForMultipleChoice,
	        MraForQuestionAnswering,
	        MraForSequenceClassification,
	        MraForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	lowerCamelCase__			 =       False
	lowerCamelCase__			 =       False
	lowerCamelCase__			 =       False
	lowerCamelCase__			 =       False

	lowerCamelCase__			 =       ()







	def      A	(						self			:     int		)			->		Union[str, Any]:
								UpperCAmelCase :       List[str] 							=		MraModelTester(self		)
								UpperCAmelCase :       Optional[int] 							=		ConfigTester(self		,			config_class=__snake_case		,			hidden_size=37		)







	def      A	(						self			:     Optional[Any]		)			->		str:
								self.config_tester.run_common_tests()







	def      A	(						self			:     Tuple		)			->		Optional[Any]:
								UpperCAmelCase :       Optional[int] 							=		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*__snake_case		)







	def      A	(						self			:     List[Any]		)			->		Optional[Any]:
								UpperCAmelCase :       Tuple 							=		self.model_tester.prepare_config_and_inputs()
								for type in ["absolute", "relative_key", "relative_key_query"]:
															UpperCAmelCase :       List[Any] 							=		type
															self.model_tester.create_and_check_model(*__snake_case		)







	def      A	(						self			:     Tuple		)			->		Dict:
								UpperCAmelCase :       List[Any] 							=		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*__snake_case		)







	def      A	(						self			:     Tuple		)			->		List[str]:
								UpperCAmelCase :       Tuple 							=		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_multiple_choice(*__snake_case		)







	def      A	(						self			:     int		)			->		Dict:
								UpperCAmelCase :       Optional[int] 							=		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*__snake_case		)







	def      A	(						self			:     Dict		)			->		Optional[int]:
								UpperCAmelCase :       Any 							=		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*__snake_case		)







	def      A	(						self			:     Any		)			->		Optional[int]:
								UpperCAmelCase :       List[str] 							=		self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*__snake_case		)







	@slow
	def      A	(						self			:     Dict		)			->		Any:
								for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															UpperCAmelCase :       str 							=		MraModel.from_pretrained(__snake_case		)
															self.assertIsNotNone(__snake_case		)







	@unittest.skip(reason='''MRA does not output attentions'''		)
	def      A	(						self			:     str		)			->		Optional[Any]:
								return

@require_torch
class 				SCREAMING_SNAKE_CASE(							unittest.TestCase			):







	"""simple docstring"""






	@slow
	def      A	(						self			:     Tuple		)			->		List[Any]:
								UpperCAmelCase :       int 							=		MraModel.from_pretrained('''uw-madison/mra-base-512-4'''		)
								UpperCAmelCase :       Optional[Any] 							=		torch.arange(256		).unsqueeze(0		)

								with torch.no_grad():
															UpperCAmelCase :       List[Any] 							=		model(__snake_case		)[0]

								UpperCAmelCase :       Optional[Any] 							=		torch.Size((1, 256, 768)		)
								self.assertEqual(output.shape		,			__snake_case		)

								UpperCAmelCase :       Any 							=		torch.tensor(
								    [[[-0.01_40, 0.08_30, -0.03_81], [0.15_46, 0.14_02, 0.02_20], [0.11_62, 0.08_51, 0.01_65]]]		)

								self.assertTrue(torch.allclose(output[:, :3, :3]		,			__snake_case		,			atol=1E-4		)		)







	@slow
	def      A	(						self			:     Optional[Any]		)			->		Any:
								UpperCAmelCase :       Optional[int] 							=		MraForMaskedLM.from_pretrained('''uw-madison/mra-base-512-4'''		)
								UpperCAmelCase :       Dict 							=		torch.arange(256		).unsqueeze(0		)

								with torch.no_grad():
															UpperCAmelCase :       List[Any] 							=		model(__snake_case		)[0]

								UpperCAmelCase :       int 							=		50265

								UpperCAmelCase :       int 							=		torch.Size((1, 256, vocab_size)		)
								self.assertEqual(output.shape		,			__snake_case		)

								UpperCAmelCase :       Union[str, Any] 							=		torch.tensor(
								    [[[9.25_95, -3.60_38, 11.88_19], [9.38_69, -3.26_93, 11.09_56], [11.85_24, -3.49_38, 13.12_10]]]		)

								self.assertTrue(torch.allclose(output[:, :3, :3]		,			__snake_case		,			atol=1E-4		)		)







	@slow
	def      A	(						self			:     str		)			->		List[Any]:
								UpperCAmelCase :       List[Any] 							=		MraForMaskedLM.from_pretrained('''uw-madison/mra-base-4096-8-d3'''		)
								UpperCAmelCase :       List[Any] 							=		torch.arange(4096		).unsqueeze(0		)

								with torch.no_grad():
															UpperCAmelCase :       Tuple 							=		model(__snake_case		)[0]

								UpperCAmelCase :       Optional[int] 							=		50265

								UpperCAmelCase :       Tuple 							=		torch.Size((1, 4096, vocab_size)		)
								self.assertEqual(output.shape		,			__snake_case		)

								UpperCAmelCase :       Optional[int] 							=		torch.tensor(
								    [[[5.47_89, -2.35_64, 7.50_64], [7.90_67, -1.33_69, 9.96_68], [9.07_12, -1.81_06, 7.03_80]]]		)

								self.assertTrue(torch.allclose(output[:, :3, :3]		,			__snake_case		,			atol=1E-4		)		)





