import argparse
import collections

import torch
from flax import traverse_util
from tax import checkpoints

from transformers import TaConfig, TaEncoderModel, TaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE="attention"			):
    '''simple docstring'''







    __UpperCamelCase  :int    						=	params[f"""{prefix}/layers_{i}/{layer_name}/key/kernel"""]
    __UpperCamelCase  :Optional[int]    						=	params[f"""{prefix}/layers_{i}/{layer_name}/out/kernel"""]
    __UpperCamelCase  :List[str]    						=	params[f"""{prefix}/layers_{i}/{layer_name}/query/kernel"""]
    __UpperCamelCase  :Union[str, Any]    						=	params[f"""{prefix}/layers_{i}/{layer_name}/value/kernel"""]
    return k, o, q, v



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE=False			):
    '''simple docstring'''







    if split_mlp_wi:
        __UpperCamelCase  :List[str]    						=	params[f"""{prefix}/layers_{i}/mlp/wi_0/kernel"""]
        __UpperCamelCase  :Dict    						=	params[f"""{prefix}/layers_{i}/mlp/wi_1/kernel"""]
        __UpperCamelCase  :Optional[Any]    						=	(wi_a, wi_a)
    else:
        __UpperCamelCase  :List[str]    						=	params[f"""{prefix}/layers_{i}/mlp/wi/kernel"""]

    __UpperCamelCase  :int    						=	params[f"""{prefix}/layers_{i}/mlp/wo/kernel"""]
    return wi, wo



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    return params[f"""{prefix}/layers_{i}/{layer_name}/scale"""]



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    *, SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    __UpperCamelCase  :Tuple    						=	traverse_util.flatten_dict(variables['''target''']			)
    __UpperCamelCase  :Optional[int]    						=	{'''/'''.join(SCREAMING_SNAKE_CASE			): v for k, v in old.items()}

    # v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
    __UpperCamelCase  :Union[str, Any]    						=	'''encoder/layers_0/mlp/wi_0/kernel''' in old
    print('''Split MLP:'''						,    SCREAMING_SNAKE_CASE			)

    __UpperCamelCase  :Optional[Any]    						=	collections.OrderedDict()

    # Shared embeddings.
    __UpperCamelCase  :Any    						=	old['''token_embedder/embedding''']

    # Encoder.
    for i in range(SCREAMING_SNAKE_CASE			):
        # Block i, layer 0 (Self Attention).
        __UpperCamelCase  :Tuple    						=	tax_layer_norm_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''encoder'''						,    '''pre_attention_layer_norm'''			)
        __UpperCamelCase  :Dict    						=	tax_attention_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''encoder'''						,    '''attention'''			)
        __UpperCamelCase  :str    						=	layer_norm
        __UpperCamelCase  :str    						=	k.T
        __UpperCamelCase  :Optional[Any]    						=	o.T
        __UpperCamelCase  :Optional[int]    						=	q.T
        __UpperCamelCase  :Any    						=	v.T

        # Block i, layer 1 (MLP).
        __UpperCamelCase  :Union[str, Any]    						=	tax_layer_norm_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''encoder'''						,    '''pre_mlp_layer_norm'''			)
        __UpperCamelCase  :Optional[int]    						=	tax_mlp_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''encoder'''						,    SCREAMING_SNAKE_CASE			)
        __UpperCamelCase  :Any    						=	layer_norm
        if split_mlp_wi:
            __UpperCamelCase  :Optional[Any]    						=	wi[0].T
            __UpperCamelCase  :int    						=	wi[1].T
        else:
            __UpperCamelCase  :Dict    						=	wi.T
        __UpperCamelCase  :str    						=	wo.T

    __UpperCamelCase  :Optional[Any]    						=	old[
        '''encoder/relpos_bias/rel_embedding'''
    ].T
    __UpperCamelCase  :Optional[int]    						=	old['''encoder/encoder_norm/scale''']

    if not is_encoder_only:
        # Decoder.
        for i in range(SCREAMING_SNAKE_CASE			):
            # Block i, layer 0 (Self Attention).
            __UpperCamelCase  :int    						=	tax_layer_norm_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''decoder'''						,    '''pre_self_attention_layer_norm'''			)
            __UpperCamelCase  :Dict    						=	tax_attention_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''decoder'''						,    '''self_attention'''			)
            __UpperCamelCase  :str    						=	layer_norm
            __UpperCamelCase  :Optional[Any]    						=	k.T
            __UpperCamelCase  :List[str]    						=	o.T
            __UpperCamelCase  :Union[str, Any]    						=	q.T
            __UpperCamelCase  :int    						=	v.T

            # Block i, layer 1 (Cross Attention).
            __UpperCamelCase  :Optional[Any]    						=	tax_layer_norm_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''decoder'''						,    '''pre_cross_attention_layer_norm'''			)
            __UpperCamelCase  :int    						=	tax_attention_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''decoder'''						,    '''encoder_decoder_attention'''			)
            __UpperCamelCase  :List[str]    						=	layer_norm
            __UpperCamelCase  :Any    						=	k.T
            __UpperCamelCase  :Any    						=	o.T
            __UpperCamelCase  :Optional[int]    						=	q.T
            __UpperCamelCase  :Union[str, Any]    						=	v.T

            # Block i, layer 2 (MLP).
            __UpperCamelCase  :Tuple    						=	tax_layer_norm_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''decoder'''						,    '''pre_mlp_layer_norm'''			)
            __UpperCamelCase  :Optional[Any]    						=	tax_mlp_lookup(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    '''decoder'''						,    SCREAMING_SNAKE_CASE			)
            __UpperCamelCase  :Optional[Any]    						=	layer_norm
            if split_mlp_wi:
                __UpperCamelCase  :int    						=	wi[0].T
                __UpperCamelCase  :Any    						=	wi[1].T
            else:
                __UpperCamelCase  :Optional[Any]    						=	wi.T
            __UpperCamelCase  :Dict    						=	wo.T

        __UpperCamelCase  :Optional[int]    						=	old['''decoder/decoder_norm/scale''']
        __UpperCamelCase  :Optional[Any]    						=	old[
            '''decoder/relpos_bias/rel_embedding'''
        ].T

        # LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
        if "decoder/logits_dense/kernel" in old:
            __UpperCamelCase  :List[str]    						=	old['''decoder/logits_dense/kernel'''].T

    return new



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    __UpperCamelCase  :Optional[Any]    						=	collections.OrderedDict([(k, torch.from_numpy(v.copy()			)) for (k, v) in converted_params.items()]			)

    # Add what is missing.
    if "encoder.embed_tokens.weight" not in state_dict:
        __UpperCamelCase  :Union[str, Any]    						=	state_dict['''shared.weight''']

    if not is_encoder_only:
        if "decoder.embed_tokens.weight" not in state_dict:
            __UpperCamelCase  :Any    						=	state_dict['''shared.weight''']

        if "lm_head.weight" not in state_dict:  # For old 1.0 models.
            print('''Using shared word embeddings as lm_head.'''			)
            __UpperCamelCase  :Optional[int]    						=	state_dict['''shared.weight''']

    return state_dict



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			):
    '''simple docstring'''







    __UpperCamelCase  :int    						=	checkpoints.load_tax_checkpoint(SCREAMING_SNAKE_CASE			)
    __UpperCamelCase  :str    						=	convert_tax_to_pytorch(SCREAMING_SNAKE_CASE						,    num_layers=config.num_layers						,    is_encoder_only=SCREAMING_SNAKE_CASE			)
    __UpperCamelCase  :Optional[int]    						=	make_state_dict(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			)
    model.load_state_dict(SCREAMING_SNAKE_CASE						,    strict=SCREAMING_SNAKE_CASE			)



def lowerCamelCase						(     SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE = False			):
    '''simple docstring'''







    __UpperCamelCase  :Union[str, Any]    						=	TaConfig.from_json_file(SCREAMING_SNAKE_CASE			)
    print(f"""Building PyTorch model from configuration: {config}"""			)
    # Non-v1.1 checkpoints could also use T5Model, but this works for all.
    # The v1.0 checkpoints will simply have an LM head that is the word embeddings.
    if is_encoder_only:
        __UpperCamelCase  :Dict    						=	TaEncoderModel(SCREAMING_SNAKE_CASE			)
    else:
        __UpperCamelCase  :str    						=	TaForConditionalGeneration(SCREAMING_SNAKE_CASE			)

    # Load weights from tf checkpoint
    load_tax_weights_in_ta(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE			)

    # Save pytorch-model
    print(f"""Save PyTorch model to {pytorch_dump_path}"""			)
    model.save_pretrained(SCREAMING_SNAKE_CASE			)

    # Verify that we can load the checkpoint.
    model.from_pretrained(SCREAMING_SNAKE_CASE			)
    print('''Done'''			)


if __name__ == "__main__":
    __lowercase											=      argparse.ArgumentParser(description='''Converts a native T5X checkpoint into a PyTorch checkpoint.''')
    # Required parameters
    parser.add_argument(
        '''--t5x_checkpoint_path''', default=None, type=str, required=True, help='''Path to the T5X checkpoint.'''
    )
    parser.add_argument(
        '''--config_file''',
        default=None,
        type=str,
        required=True,
        help='''The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.''',
    )
    parser.add_argument(
        '''--pytorch_dump_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
    )
    parser.add_argument(
        '''--is_encoder_only''', action='''store_true''', help='''Check if the model is encoder-decoder model''', default=False
    )
    __lowercase											=      parser.parse_args()
    convert_tax_checkpoint_to_pytorch(
        args.tax_checkpoint_path, args.config_file, args.pytorch_dump_path, args.is_encoder_only
    )






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






    '''simple docstring'''






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







    def     UpperCamelCase__     (			self) ->			str:
        __UpperCamelCase  :int    						=	ids_tensor([self.batch_size, self.seq_length]				,				self.vocab_size)

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

        __UpperCamelCase  :Union[str, Any]    						=	None
        if self.use_token_type_ids:
            __UpperCamelCase  :Optional[int]    						=	ids_tensor([self.batch_size, self.seq_length]				,				self.type_vocab_size)

        __UpperCamelCase  :List[str]    						=	None
        __UpperCamelCase  :Tuple    						=	None
        __UpperCamelCase  :Union[str, Any]    						=	None
        if self.use_labels:
            __UpperCamelCase  :Dict    						=	ids_tensor([self.batch_size]				,				self.type_sequence_label_size)
            __UpperCamelCase  :Optional[int]    						=	ids_tensor([self.batch_size, self.seq_length]				,				self.num_labels)
            __UpperCamelCase  :List[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     UpperCamelCase__     (			self) ->			Union[str, 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=__lowercase				,				initializer_range=self.initializer_range				,				)







    def     UpperCamelCase__     (			self) ->			List[Any]:
        __UpperCamelCase  :Dict    						=	self.get_config()
        __UpperCamelCase  :List[str]    						=	300
        return config







    def     UpperCamelCase__     (			self) ->			int:
        (
            (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, 
        )  :Tuple    						=	self.prepare_config_and_inputs()

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







    def     UpperCamelCase__     (			self				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				) ->			List[Any]:
        __UpperCamelCase  :Tuple    						=	True
        __UpperCamelCase  :Dict    						=	MraModel(__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase  :int    						=	model(
            __lowercase				,				attention_mask=__lowercase				,				token_type_ids=__lowercase				,				encoder_hidden_states=__lowercase				,				encoder_attention_mask=__lowercase				,				)
        __UpperCamelCase  :Any    						=	model(
            __lowercase				,				attention_mask=__lowercase				,				token_type_ids=__lowercase				,				encoder_hidden_states=__lowercase				,				)
        __UpperCamelCase  :str    						=	model(__lowercase				,				attention_mask=__lowercase				,				token_type_ids=__lowercase)
        self.parent.assertEqual(result.last_hidden_state.shape				,				(self.batch_size, self.seq_length, self.hidden_size))







    def     UpperCamelCase__     (			self				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase) ->			Union[str, Any]:
        __UpperCamelCase  :Dict    						=	MraForMaskedLM(config=__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase  :Optional[Any]    						=	model(__lowercase				,				attention_mask=__lowercase				,				token_type_ids=__lowercase				,				labels=__lowercase)
        self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.seq_length, self.vocab_size))







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







    def     UpperCamelCase__     (			self				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase				,				__lowercase) ->			Optional[int]:
        __UpperCamelCase  :str    						=	self.num_labels
        __UpperCamelCase  :Optional[int]    						=	MraForTokenClassification(config=__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase  :Optional[Any]    						=	model(__lowercase				,				attention_mask=__lowercase				,				token_type_ids=__lowercase				,				labels=__lowercase)
        self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.seq_length, self.num_labels))







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







    def     UpperCamelCase__     (			self) ->			Optional[int]:
        __UpperCamelCase  :str    						=	self.prepare_config_and_inputs()
        (
            (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, (
            __UpperCamelCase
        )						, 
        )  :str    						=	config_and_inputs
        __UpperCamelCase  :List[str]    						=	{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
        return config, inputs_dict


@require_torch
class 			lowerCamelCase_  (     UpperCAmelCase_		,			unittest.TestCase				):






    '''simple docstring'''






    a__			:		List[str]      			=   (
        (
            MraModel,
            MraForMaskedLM,
            MraForMultipleChoice,
            MraForQuestionAnswering,
            MraForSequenceClassification,
            MraForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    a__			:		Optional[int]      			=   False
    a__			:		Optional[int]      			=   False
    a__			:		str      			=   False
    a__			:		Optional[Any]      			=   False

    a__			:		Union[str, Any]      			=   ()







    def     UpperCamelCase__     (			self) ->			List[str]:
        __UpperCamelCase  :str    						=	MraModelTester(self)
        __UpperCamelCase  :Tuple    						=	ConfigTester(self				,				config_class=__lowercase				,				hidden_size=37)







    def     UpperCamelCase__     (			self) ->			Dict:
        self.config_tester.run_common_tests()







    def     UpperCamelCase__     (			self) ->			Union[str, Any]:
        __UpperCamelCase  :List[str]    						=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowercase)







    def     UpperCamelCase__     (			self) ->			List[Any]:
        __UpperCamelCase  :Union[str, Any]    						=	self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
            __UpperCamelCase  :Optional[int]    						=	type
            self.model_tester.create_and_check_model(*__lowercase)







    def     UpperCamelCase__     (			self) ->			List[Any]:
        __UpperCamelCase  :Union[str, Any]    						=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*__lowercase)







    def     UpperCamelCase__     (			self) ->			Any:
        __UpperCamelCase  :Dict    						=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*__lowercase)







    def     UpperCamelCase__     (			self) ->			Dict:
        __UpperCamelCase  :Union[str, Any]    						=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*__lowercase)







    def     UpperCamelCase__     (			self) ->			Tuple:
        __UpperCamelCase  :Optional[int]    						=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*__lowercase)







    def     UpperCamelCase__     (			self) ->			Any:
        __UpperCamelCase  :Dict    						=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*__lowercase)







    @slow
    def     UpperCamelCase__     (			self) ->			Dict:
        for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            __UpperCamelCase  :List[Any]    						=	MraModel.from_pretrained(__lowercase)
            self.assertIsNotNone(__lowercase)







    @unittest.skip(reason='''MRA does not output attentions''')
    def     UpperCamelCase__     (			self) ->			Any:
        return


@require_torch
class 			lowerCamelCase_  (     unittest.TestCase				):






    '''simple docstring'''






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

        with torch.no_grad():
            __UpperCamelCase  :List[Any]    						=	model(__lowercase)[0]

        __UpperCamelCase  :Dict    						=	torch.Size((1, 256, 768))
        self.assertEqual(output.shape				,				__lowercase)

        __UpperCamelCase  :Optional[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]				,				__lowercase				,				atol=1E-4))







    @slow
    def     UpperCamelCase__     (			self) ->			Any:
        __UpperCamelCase  :str    						=	MraForMaskedLM.from_pretrained('''uw-madison/mra-base-512-4''')
        __UpperCamelCase  :str    						=	torch.arange(256).unsqueeze(0)

        with torch.no_grad():
            __UpperCamelCase  :Optional[Any]    						=	model(__lowercase)[0]

        __UpperCamelCase  :List[str]    						=	50_265

        __UpperCamelCase  :List[Any]    						=	torch.Size((1, 256, vocab_size))
        self.assertEqual(output.shape				,				__lowercase)

        __UpperCamelCase  :Optional[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]				,				__lowercase				,				atol=1E-4))







    @slow
    def     UpperCamelCase__     (			self) ->			Optional[int]:
        __UpperCamelCase  :Dict    						=	MraForMaskedLM.from_pretrained('''uw-madison/mra-base-4096-8-d3''')
        __UpperCamelCase  :Optional[int]    						=	torch.arange(4_096).unsqueeze(0)

        with torch.no_grad():
            __UpperCamelCase  :Tuple    						=	model(__lowercase)[0]

        __UpperCamelCase  :Optional[int]    						=	50_265

        __UpperCamelCase  :Optional[Any]    						=	torch.Size((1, 4_096, vocab_size))
        self.assertEqual(output.shape				,				__lowercase)

        __UpperCamelCase  :List[str]    						=	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]				,				__lowercase				,				atol=1E-4))





