



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


a      :  List[Any]			     =			logging.get_logger(__name__)

# General docstring
a      :  Tuple			     =			'RegNetConfig'

# Base docstring
a      :  Dict			     =			'facebook/regnet-y-040'
a      :  Tuple			     =			[1, 1_088, 7, 7]

# Image classification docstring
a      :  str			     =			'facebook/regnet-y-040'
a      :  Tuple			     =			'tabby, tabby cat'

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







class 				_a				(   nn.Module					):

			def __init__(self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_ = 3,   SCREAMING_SNAKE_CASE_ = 1,   SCREAMING_SNAKE_CASE_ = 1,   SCREAMING_SNAKE_CASE_ = "relu",   )       ->   Union[str, Any]:
						super().__init__()
						UpperCAmelCase_:		Optional[Any]							      =     nn.Convad(
						    SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   kernel_size=SCREAMING_SNAKE_CASE_,   stride=SCREAMING_SNAKE_CASE_,   padding=kernel_size // 2,   groups=SCREAMING_SNAKE_CASE_,   bias=SCREAMING_SNAKE_CASE_,   )
						UpperCAmelCase_:		List[Any]							      =     nn.BatchNormad(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		List[Any]							      =     ACTaFN[activation] if activation is not None else nn.Identity()

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_					)       ->   Any:
						UpperCAmelCase_:		List[Any]							      =     self.convolution(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Tuple							      =     self.normalization(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Tuple							      =     self.activation(SCREAMING_SNAKE_CASE_					)
						return hidden_state







class 				_a				(   nn.Module					):

			def __init__(self,   SCREAMING_SNAKE_CASE_					)       ->   List[Any]:
						super().__init__()
						UpperCAmelCase_:		Union[str, Any]							      =     RegNetConvLayer(
						    config.num_channels,   config.embedding_size,   kernel_size=3,   stride=2,   activation=config.hidden_act					)
						UpperCAmelCase_:		Dict							      =     config.num_channels

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_					)       ->   Dict:
						UpperCAmelCase_:		Optional[int]							      =     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."""					)
						UpperCAmelCase_:		int							      =     self.embedder(SCREAMING_SNAKE_CASE_					)
						return hidden_state







class 				_a				(   nn.Module					):

			def __init__(self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_ = 2					)       ->   str:
						super().__init__()
						UpperCAmelCase_:		str							      =     nn.Convad(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   kernel_size=1,   stride=SCREAMING_SNAKE_CASE_,   bias=SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Optional[Any]							      =     nn.BatchNormad(SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_					)       ->   Tensor:
						UpperCAmelCase_:		Optional[Any]							      =     self.convolution(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Union[str, Any]							      =     self.normalization(SCREAMING_SNAKE_CASE_					)
						return hidden_state







class 				_a				(   nn.Module					):

			def __init__(self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Any:
						super().__init__()

						UpperCAmelCase_:		Dict							      =     nn.AdaptiveAvgPoolad((1, 1)					)
						UpperCAmelCase_:		Dict							      =     nn.Sequential(
						    nn.Convad(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   kernel_size=1					),   nn.ReLU(),   nn.Convad(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   kernel_size=1					),   nn.Sigmoid(),   )

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_					)       ->   List[str]:
						# b c h w -> b c 1 1
						UpperCAmelCase_:		int							      =     self.pooler(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Dict							      =     self.attention(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Any							      =     hidden_state * attention
						return hidden_state







class 				_a				(   nn.Module					):

			def __init__(self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_ = 1					)       ->   Tuple:
						super().__init__()
						UpperCAmelCase_:		List[str]							      =     in_channels != out_channels or stride != 1
						UpperCAmelCase_:		List[str]							      =     max(1,   out_channels // config.groups_width					)
						UpperCAmelCase_:		int							      =     (
						    RegNetShortCut(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   stride=SCREAMING_SNAKE_CASE_					) if should_apply_shortcut else nn.Identity()
						)
						UpperCAmelCase_:		List[Any]							      =     nn.Sequential(
						    RegNetConvLayer(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   kernel_size=1,   activation=config.hidden_act					),   RegNetConvLayer(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   stride=SCREAMING_SNAKE_CASE_,   groups=SCREAMING_SNAKE_CASE_,   activation=config.hidden_act					),   RegNetConvLayer(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   kernel_size=1,   activation=SCREAMING_SNAKE_CASE_					),   )
						UpperCAmelCase_:		str							      =     ACTaFN[config.hidden_act]

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_					)       ->   List[Any]:
						UpperCAmelCase_:		List[str]							      =     hidden_state
						UpperCAmelCase_:		Dict							      =     self.layer(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Tuple							      =     self.shortcut(SCREAMING_SNAKE_CASE_					)
						hidden_state += residual
						UpperCAmelCase_:		List[str]							      =     self.activation(SCREAMING_SNAKE_CASE_					)
						return hidden_state







class 				_a				(   nn.Module					):

			def __init__(self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_ = 1					)       ->   str:
						super().__init__()
						UpperCAmelCase_:		Optional[int]							      =     in_channels != out_channels or stride != 1
						UpperCAmelCase_:		List[str]							      =     max(1,   out_channels // config.groups_width					)
						UpperCAmelCase_:		List[Any]							      =     (
						    RegNetShortCut(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   stride=SCREAMING_SNAKE_CASE_					) if should_apply_shortcut else nn.Identity()
						)
						UpperCAmelCase_:		Any							      =     nn.Sequential(
						    RegNetConvLayer(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   kernel_size=1,   activation=config.hidden_act					),   RegNetConvLayer(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   stride=SCREAMING_SNAKE_CASE_,   groups=SCREAMING_SNAKE_CASE_,   activation=config.hidden_act					),   RegNetSELayer(SCREAMING_SNAKE_CASE_,   reduced_channels=int(round(in_channels / 4					)					)					),   RegNetConvLayer(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   kernel_size=1,   activation=SCREAMING_SNAKE_CASE_					),   )
						UpperCAmelCase_:		Union[str, Any]							      =     ACTaFN[config.hidden_act]

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_					)       ->   Optional[int]:
						UpperCAmelCase_:		Any							      =     hidden_state
						UpperCAmelCase_:		List[Any]							      =     self.layer(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		List[str]							      =     self.shortcut(SCREAMING_SNAKE_CASE_					)
						hidden_state += residual
						UpperCAmelCase_:		List[Any]							      =     self.activation(SCREAMING_SNAKE_CASE_					)
						return hidden_state







class 				_a				(   nn.Module					):

			def __init__(self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_ = 2,   SCREAMING_SNAKE_CASE_ = 2,   )       ->   Any:
						super().__init__()

						UpperCAmelCase_:		List[str]							      =     RegNetXLayer if config.layer_type == """x""" else RegNetYLayer

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

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_					)       ->   Optional[Any]:
						UpperCAmelCase_:		Optional[Any]							      =     self.layers(SCREAMING_SNAKE_CASE_					)
						return hidden_state







class 				_a				(   nn.Module					):

			def __init__(self,   SCREAMING_SNAKE_CASE_					)       ->   Any:
						super().__init__()
						UpperCAmelCase_:		Optional[Any]							      =     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(
						        SCREAMING_SNAKE_CASE_,   config.embedding_size,   config.hidden_sizes[0],   stride=2 if config.downsample_in_first_stage else 1,   depth=config.depths[0],   )					)
						UpperCAmelCase_:		List[str]							      =     zip(config.hidden_sizes,   config.hidden_sizes[1:]					)
						for (in_channels, out_channels), depth in zip(SCREAMING_SNAKE_CASE_,   config.depths[1:]					):
									self.stages.append(RegNetStage(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   depth=SCREAMING_SNAKE_CASE_					)					)

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_ = False,   SCREAMING_SNAKE_CASE_ = True					)       ->   BaseModelOutputWithNoAttention:
						UpperCAmelCase_:		List[str]							      =     () if output_hidden_states else None

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

									UpperCAmelCase_:		Any							      =     stage_module(SCREAMING_SNAKE_CASE_					)

						if output_hidden_states:
									UpperCAmelCase_:		Optional[Any]							      =     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=SCREAMING_SNAKE_CASE_,   hidden_states=SCREAMING_SNAKE_CASE_					)







class 				_a				(   _lowerCAmelCase					):
			A      		=	RegNetConfig
			A      		=	'''regnet'''
			A      		=	'''pixel_values'''
			A      		=	True

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

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_=False					)       ->   List[str]:
						if isinstance(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					):
									UpperCAmelCase_:		Tuple							      =     value






a      :  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'

a      :  int			     =			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.''' ,		_lowerCAmelCase ,		)
# Copied from transformers.models.resnet.modeling_resnet.ResNetModel with RESNET->REGNET,ResNet->RegNet
class 				_a				(   _lowerCAmelCase					):

			def __init__(self,   SCREAMING_SNAKE_CASE_					)       ->   Tuple:
						super().__init__(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		List[Any]							      =     config
						UpperCAmelCase_:		int							      =     RegNetEmbeddings(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Dict							      =     RegNetEncoder(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Tuple							      =     nn.AdaptiveAvgPoolad((1, 1)					)
						# Initialize weights and apply final processing
						self.post_init()

			@add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE_					)
			@add_code_sample_docstrings(
			    checkpoint=_CHECKPOINT_FOR_DOC,   output_type=SCREAMING_SNAKE_CASE_,   config_class=_CONFIG_FOR_DOC,   modality="""vision""",   expected_output=_EXPECTED_OUTPUT_SHAPE,   )
			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_ = None,   SCREAMING_SNAKE_CASE_ = None					)       ->   BaseModelOutputWithPoolingAndNoAttention:
						UpperCAmelCase_:		Optional[Any]							      =     (
						    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
						)
						UpperCAmelCase_:		Optional[Any]							      =     return_dict if return_dict is not None else self.config.use_return_dict

						UpperCAmelCase_:		List[Any]							      =     self.embedder(SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Optional[Any]							      =     self.encoder(
						    SCREAMING_SNAKE_CASE_,   output_hidden_states=SCREAMING_SNAKE_CASE_,   return_dict=SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		int							      =     encoder_outputs[0]

						UpperCAmelCase_:		List[Any]							      =     self.pooler(SCREAMING_SNAKE_CASE_					)

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

						return BaseModelOutputWithPoolingAndNoAttention(
						    last_hidden_state=SCREAMING_SNAKE_CASE_,   pooler_output=SCREAMING_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.
    ''' ,		_lowerCAmelCase ,		)
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class 				_a				(   _lowerCAmelCase					):

			def __init__(self,   SCREAMING_SNAKE_CASE_					)       ->   List[Any]:
						super().__init__(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Dict							      =     config.num_labels
						UpperCAmelCase_:		Union[str, Any]							      =     RegNetModel(SCREAMING_SNAKE_CASE_					)
						# classification head
						UpperCAmelCase_:		Optional[int]							      =     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(SCREAMING_SNAKE_CASE_					)
			@add_code_sample_docstrings(
			    checkpoint=_IMAGE_CLASS_CHECKPOINT,   output_type=SCREAMING_SNAKE_CASE_,   config_class=_CONFIG_FOR_DOC,   expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,   )
			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_ = None,   SCREAMING_SNAKE_CASE_ = None,   SCREAMING_SNAKE_CASE_ = None,   SCREAMING_SNAKE_CASE_ = None,   )       ->   ImageClassifierOutputWithNoAttention:
						UpperCAmelCase_:		Union[str, Any]							      =     return_dict if return_dict is not None else self.config.use_return_dict

						UpperCAmelCase_:		Optional[int]							      =     self.regnet(SCREAMING_SNAKE_CASE_,   output_hidden_states=SCREAMING_SNAKE_CASE_,   return_dict=SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Any							      =     outputs.pooler_output if return_dict else outputs[1]

						UpperCAmelCase_:		int							      =     self.classifier(SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Union[str, Any]							      =     None

						if labels is not None:
									if self.config.problem_type is None:
												if self.num_labels == 1:
															UpperCAmelCase_:		Tuple							      =     """regression"""
												elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
															UpperCAmelCase_:		Dict							      =     """single_label_classification"""
												else:
															UpperCAmelCase_:		List[Any]							      =     """multi_label_classification"""
									if self.config.problem_type == "regression":
												UpperCAmelCase_:		Dict							      =     MSELoss()
												if self.num_labels == 1:
															UpperCAmelCase_:		Tuple							      =     loss_fct(logits.squeeze(),   labels.squeeze()					)
												else:
															UpperCAmelCase_:		List[str]							      =     loss_fct(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)
									elif self.config.problem_type == "single_label_classification":
												UpperCAmelCase_:		List[Any]							      =     CrossEntropyLoss()
												UpperCAmelCase_:		List[Any]							      =     loss_fct(logits.view(-1,   self.num_labels					),   labels.view(-1					)					)
									elif self.config.problem_type == "multi_label_classification":
												UpperCAmelCase_:		List[str]							      =     BCEWithLogitsLoss()
												UpperCAmelCase_:		List[Any]							      =     loss_fct(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)

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

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











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

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

			def      __snake_case       (self					)       ->   Optional[Any]:
						UpperCAmelCase_:		Tuple							      =     ids_tensor([self.batch_size, self.seq_length],   self.vocab_size					)

						UpperCAmelCase_:		List[str]							      =     None
						if self.use_input_mask:
									UpperCAmelCase_:		Union[str, Any]							      =     random_attention_mask([self.batch_size, self.seq_length]					)

						UpperCAmelCase_:		List[str]							      =     None
						if self.use_token_type_ids:
									UpperCAmelCase_:		str							      =     ids_tensor([self.batch_size, self.seq_length],   self.type_vocab_size					)

						UpperCAmelCase_:		str							      =     None
						UpperCAmelCase_:		List[Any]							      =     None
						UpperCAmelCase_:		Tuple							      =     None
						if self.use_labels:
									UpperCAmelCase_:		Tuple							      =     ids_tensor([self.batch_size],   self.type_sequence_label_size					)
									UpperCAmelCase_:		str							      =     ids_tensor([self.batch_size, self.seq_length],   self.num_labels					)
									UpperCAmelCase_:		Tuple							      =     ids_tensor([self.batch_size],   self.num_choices					)

						UpperCAmelCase_:		List[Any]							      =     self.get_config()

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

			def      __snake_case       (self					)       ->   List[Any]:
						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=SCREAMING_SNAKE_CASE_,   initializer_range=self.initializer_range,   )

			def      __snake_case       (self					)       ->   Optional[int]:
						UpperCAmelCase_:		Optional[Any]							      =     self.get_config()
						UpperCAmelCase_:		Dict							      =     300
						return config

			def      __snake_case       (self					)       ->   List[Any]:
						(
						    (
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							
						):		Any							      =     self.prepare_config_and_inputs()

						UpperCAmelCase_:		List[str]							      =     True
						UpperCAmelCase_:		Any							      =     floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
						UpperCAmelCase_:		List[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      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Optional[Any]:
						UpperCAmelCase_:		Optional[int]							      =     MraModel(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		Optional[Any]							      =     model(SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Tuple							      =     model(SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Tuple							      =     model(SCREAMING_SNAKE_CASE_					)
						self.parent.assertEqual(result.last_hidden_state.shape,   (self.batch_size, self.seq_length, self.hidden_size)					)

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

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

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   List[str]:
						UpperCAmelCase_:		Any							      =     MraForQuestionAnswering(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		str							      =     model(
						    SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_,   start_positions=SCREAMING_SNAKE_CASE_,   end_positions=SCREAMING_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      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Optional[int]:
						UpperCAmelCase_:		List[str]							      =     self.num_labels
						UpperCAmelCase_:		Any							      =     MraForSequenceClassification(SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		List[str]							      =     model(SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_,   labels=SCREAMING_SNAKE_CASE_					)
						self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.num_labels)					)

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

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

			def      __snake_case       (self					)       ->   Union[str, Any]:
						UpperCAmelCase_:		int							      =     self.prepare_config_and_inputs()
						(
						    (
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							
						):		Any							      =     config_and_inputs
						UpperCAmelCase_:		Dict							      =     {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
						return config, inputs_dict







@require_torch
class 				_a				(   _lowerCAmelCase ,		unittest.TestCase					):
			A      		=	(
			    (
			        MraModel,
			        MraForMaskedLM,
			        MraForMultipleChoice,
			        MraForQuestionAnswering,
			        MraForSequenceClassification,
			        MraForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			A      		=	False
			A      		=	False
			A      		=	False
			A      		=	False

			A      		=	()

			def      __snake_case       (self					)       ->   Tuple:
						UpperCAmelCase_:		Dict							      =     MraModelTester(self					)
						UpperCAmelCase_:		List[Any]							      =     ConfigTester(self,   config_class=SCREAMING_SNAKE_CASE_,   hidden_size=37					)

			def      __snake_case       (self					)       ->   Tuple:
						self.config_tester.run_common_tests()

			def      __snake_case       (self					)       ->   str:
						UpperCAmelCase_:		int							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Optional[int]:
						UpperCAmelCase_:		List[str]							      =     self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
									UpperCAmelCase_:		str							      =     type
									self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Union[str, Any]:
						UpperCAmelCase_:		str							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   List[str]:
						UpperCAmelCase_:		Optional[Any]							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Dict:
						UpperCAmelCase_:		Union[str, Any]							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Tuple:
						UpperCAmelCase_:		Union[str, Any]							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Any:
						UpperCAmelCase_:		str							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_					)

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

			@unittest.skip(reason="""MRA does not output attentions"""					)
			def      __snake_case       (self					)       ->   List[str]:
						return







@require_torch
class 				_a				(   unittest.TestCase					):

			@slow
			def      __snake_case       (self					)       ->   Tuple:
						UpperCAmelCase_:		Union[str, Any]							      =     MraModel.from_pretrained("""uw-madison/mra-base-512-4"""					)
						UpperCAmelCase_:		Optional[Any]							      =     torch.arange(256					).unsqueeze(0					)

						with torch.no_grad():
									UpperCAmelCase_:		Optional[Any]							      =     model(SCREAMING_SNAKE_CASE_					)[0]

						UpperCAmelCase_:		Optional[int]							      =     torch.Size((1, 256, 768)					)
						self.assertEqual(output.shape,   SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		int							      =     torch.tensor(
						    [[[-0.0_1_4_0, 0.0_8_3_0, -0.0_3_8_1], [0.1_5_4_6, 0.1_4_0_2, 0.0_2_2_0], [0.1_1_6_2, 0.0_8_5_1, 0.0_1_6_5]]]					)

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

			@slow
			def      __snake_case       (self					)       ->   Any:
						UpperCAmelCase_:		List[str]							      =     MraForMaskedLM.from_pretrained("""uw-madison/mra-base-512-4"""					)
						UpperCAmelCase_:		Dict							      =     torch.arange(256					).unsqueeze(0					)

						with torch.no_grad():
									UpperCAmelCase_:		List[str]							      =     model(SCREAMING_SNAKE_CASE_					)[0]

						UpperCAmelCase_:		Any							      =     50265

						UpperCAmelCase_:		Any							      =     torch.Size((1, 256, vocab_size)					)
						self.assertEqual(output.shape,   SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Dict							      =     torch.tensor(
						    [[[9.2_5_9_5, -3.6_0_3_8, 1_1.8_8_1_9], [9.3_8_6_9, -3.2_6_9_3, 1_1.0_9_5_6], [1_1.8_5_2_4, -3.4_9_3_8, 1_3.1_2_1_0]]]					)

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

			@slow
			def      __snake_case       (self					)       ->   Optional[int]:
						UpperCAmelCase_:		Optional[int]							      =     MraForMaskedLM.from_pretrained("""uw-madison/mra-base-4096-8-d3"""					)
						UpperCAmelCase_:		List[Any]							      =     torch.arange(4096					).unsqueeze(0					)

						with torch.no_grad():
									UpperCAmelCase_:		List[Any]							      =     model(SCREAMING_SNAKE_CASE_					)[0]

						UpperCAmelCase_:		Dict							      =     50265

						UpperCAmelCase_:		int							      =     torch.Size((1, 4096, vocab_size)					)
						self.assertEqual(output.shape,   SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		List[Any]							      =     torch.tensor(
						    [[[5.4_7_8_9, -2.3_5_6_4, 7.5_0_6_4], [7.9_0_6_7, -1.3_3_6_9, 9.9_6_6_8], [9.0_7_1_2, -1.8_1_0_6, 7.0_3_8_0]]]					)

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






