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

     """simple docstring"""

     def __init__(      self							,		lowerCAmelCase_							,		lowerCAmelCase_=2							,		lowerCAmelCase_=8							,		lowerCAmelCase_=True							,		lowerCAmelCase_=True							,		lowerCAmelCase_=True							,		lowerCAmelCase_=True							,		lowerCAmelCase_=99							,		lowerCAmelCase_=16							,		lowerCAmelCase_=5							,		lowerCAmelCase_=2							,		lowerCAmelCase_=36							,		lowerCAmelCase_="gelu"							,		lowerCAmelCase_=0.0							,		lowerCAmelCase_=0.0							,		lowerCAmelCase_=5_12							,		lowerCAmelCase_=16							,		lowerCAmelCase_=2							,		lowerCAmelCase_=0.02							,		lowerCAmelCase_=3							,		lowerCAmelCase_=4							,		lowerCAmelCase_=None							,		):

       '''simple docstring'''

       a_      :	Dict      =							parent
       a_      :	int      =							batch_size
       a_      :	Union[str, Any]      =							seq_length
       a_      :	List[str]      =							is_training
       a_      :	Union[str, Any]      =							use_input_mask
       a_      :	Optional[int]      =							use_token_type_ids
       a_      :	Dict      =							use_labels
       a_      :	List[str]      =							vocab_size
       a_      :	Tuple      =							hidden_size
       a_      :	Dict      =							num_hidden_layers
       a_      :	Union[str, Any]      =							num_attention_heads
       a_      :	str      =							intermediate_size
       a_      :	int      =							hidden_act
       a_      :	str      =							hidden_dropout_prob
       a_      :	Union[str, Any]      =							attention_probs_dropout_prob
       a_      :	str      =							max_position_embeddings
       a_      :	Any      =							type_vocab_size
       a_      :	Optional[int]      =							type_sequence_label_size
       a_      :	Any      =							initializer_range
       a_      :	List[str]      =							num_labels
       a_      :	int      =							num_choices
       a_      :	Dict      =							scope


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	int      =							ids_tensor([self.batch_size, self.seq_length]							,		self.vocab_size )

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

       a_      :	List[Any]      =							None
       if self.use_token_type_ids:
         a_      :	Tuple      =							ids_tensor([self.batch_size, self.seq_length]							,		self.type_vocab_size )

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

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

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


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       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=_lowerCAmelCase							,		initializer_range=self.initializer_range							,		)


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	Optional[Any]      =							self.get_config()
       a_      :	Any      =							3_00
       return config


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       (
           (
           a_						
       )      , (
           a_						
       )      , (
           a_						
       )      , (
           a_						
       )      , (
           a_						
       )      , (
           a_						
       )      , (
           a_						
       )      , 						
       )      :	Dict      =							self.prepare_config_and_inputs()

       a_      :	int      =							True
       a_      :	Tuple      =							floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
       a_      :	int      =							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 							_lowerCAmelCase  (      self							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_ ):

       '''simple docstring'''

       a_      :	Optional[int]      =							MraModel(config=_lowerCAmelCase )
       model.to(_lowerCAmelCase )
       model.eval()
       a_      :	int      =							model(_lowerCAmelCase							,		attention_mask=_lowerCAmelCase							,		token_type_ids=_lowerCAmelCase )
       a_      :	List[str]      =							model(_lowerCAmelCase							,		token_type_ids=_lowerCAmelCase )
       a_      :	Optional[int]      =							model(_lowerCAmelCase )
       self.parent.assertEqual(result.last_hidden_state.shape							,		(self.batch_size, self.seq_length, self.hidden_size) )


     def 							_lowerCAmelCase  (      self							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		):

       '''simple docstring'''

       a_      :	Any      =							True
       a_      :	Optional[int]      =							MraModel(_lowerCAmelCase )
       model.to(_lowerCAmelCase )
       model.eval()
       a_      :	Any      =							model(
           _lowerCAmelCase							,		attention_mask=_lowerCAmelCase							,		token_type_ids=_lowerCAmelCase							,		encoder_hidden_states=_lowerCAmelCase							,		encoder_attention_mask=_lowerCAmelCase							,		)
       a_      :	List[str]      =							model(
           _lowerCAmelCase							,		attention_mask=_lowerCAmelCase							,		token_type_ids=_lowerCAmelCase							,		encoder_hidden_states=_lowerCAmelCase							,		)
       a_      :	Dict      =							model(_lowerCAmelCase							,		attention_mask=_lowerCAmelCase							,		token_type_ids=_lowerCAmelCase )
       self.parent.assertEqual(result.last_hidden_state.shape							,		(self.batch_size, self.seq_length, self.hidden_size) )


     def 							_lowerCAmelCase  (      self							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_ ):

       '''simple docstring'''

       a_      :	int      =							MraForMaskedLM(config=_lowerCAmelCase )
       model.to(_lowerCAmelCase )
       model.eval()
       a_      :	Any      =							model(_lowerCAmelCase							,		attention_mask=_lowerCAmelCase							,		token_type_ids=_lowerCAmelCase							,		labels=_lowerCAmelCase )
       self.parent.assertEqual(result.logits.shape							,		(self.batch_size, self.seq_length, self.vocab_size) )


     def 							_lowerCAmelCase  (      self							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_ ):

       '''simple docstring'''

       a_      :	Optional[int]      =							MraForQuestionAnswering(config=_lowerCAmelCase )
       model.to(_lowerCAmelCase )
       model.eval()
       a_      :	str      =							model(
           _lowerCAmelCase							,		attention_mask=_lowerCAmelCase							,		token_type_ids=_lowerCAmelCase							,		start_positions=_lowerCAmelCase							,		end_positions=_lowerCAmelCase							,		)
       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 							_lowerCAmelCase  (      self							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_ ):

       '''simple docstring'''

       a_      :	Optional[Any]      =							self.num_labels
       a_      :	Optional[int]      =							MraForSequenceClassification(_lowerCAmelCase )
       model.to(_lowerCAmelCase )
       model.eval()
       a_      :	Dict      =							model(_lowerCAmelCase							,		attention_mask=_lowerCAmelCase							,		token_type_ids=_lowerCAmelCase							,		labels=_lowerCAmelCase )
       self.parent.assertEqual(result.logits.shape							,		(self.batch_size, self.num_labels) )


     def 							_lowerCAmelCase  (      self							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_ ):

       '''simple docstring'''

       a_      :	Any      =							self.num_labels
       a_      :	List[str]      =							MraForTokenClassification(config=_lowerCAmelCase )
       model.to(_lowerCAmelCase )
       model.eval()
       a_      :	Union[str, Any]      =							model(_lowerCAmelCase							,		attention_mask=_lowerCAmelCase							,		token_type_ids=_lowerCAmelCase							,		labels=_lowerCAmelCase )
       self.parent.assertEqual(result.logits.shape							,		(self.batch_size, self.seq_length, self.num_labels) )


     def 							_lowerCAmelCase  (      self							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_ ):

       '''simple docstring'''

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






     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	Dict      =							self.prepare_config_and_inputs()
       (
           (
           a_						
       )      , (
           a_						
       )      , (
           a_						
       )      , (
           a_						
       )      , (
           a_						
       )      , (
           a_						
       )      , (
           a_						
       )      , 						
       )      :	Dict      =							config_and_inputs
       a_      :	int      =							{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
       return config, inputs_dict







@require_torch
class  _UpperCAmelCase ( _lowerCAmelCase	,unittest.TestCase    ):

     """simple docstring"""

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

     a_    	=					()


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	Dict      =							MraModelTester(self )
       a_      :	int      =							ConfigTester(self							,		config_class=_lowerCAmelCase							,		hidden_size=37 )


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       self.config_tester.run_common_tests()


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	Optional[int]      =							self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_lowerCAmelCase )


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	Union[str, Any]      =							self.model_tester.prepare_config_and_inputs()
       for type in ["absolute", "relative_key", "relative_key_query"]:
         a_      :	List[str]      =							type
         self.model_tester.create_and_check_model(*_lowerCAmelCase )


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	Dict      =							self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*_lowerCAmelCase )


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	List[str]      =							self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_multiple_choice(*_lowerCAmelCase )


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	List[str]      =							self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*_lowerCAmelCase )


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	List[Any]      =							self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(*_lowerCAmelCase )


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	List[str]      =							self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*_lowerCAmelCase )


     @slow
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         a_      :	Tuple      =							MraModel.from_pretrained(_lowerCAmelCase )
         self.assertIsNotNone(_lowerCAmelCase )






     @unittest.skip(reason="""MRA does not output attentions""" )
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       return







@require_torch
class  _UpperCAmelCase ( unittest.TestCase    ):

     """simple docstring"""

     @slow
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	Dict      =							MraModel.from_pretrained("""uw-madison/mra-base-512-4""" )
       a_      :	Union[str, Any]      =							torch.arange(2_56 ).unsqueeze(0 )

       with torch.no_grad():
         a_      :	int      =							model(_lowerCAmelCase )[0]

       a_      :	Any      =							torch.Size((1, 2_56, 7_68) )
       self.assertEqual(output.shape							,		_lowerCAmelCase )

       a_      :	List[Any]      =							torch.tensor(
           [[[-0.0140, 0.0830, -0.0381], [0.1546, 0.1402, 0.0220], [0.1162, 0.0851, 0.0165]]] )

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


     @slow
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	int      =							MraForMaskedLM.from_pretrained("""uw-madison/mra-base-512-4""" )
       a_      :	List[str]      =							torch.arange(2_56 ).unsqueeze(0 )

       with torch.no_grad():
         a_      :	List[Any]      =							model(_lowerCAmelCase )[0]

       a_      :	int      =							5_02_65

       a_      :	int      =							torch.Size((1, 2_56, vocab_size) )
       self.assertEqual(output.shape							,		_lowerCAmelCase )

       a_      :	Dict      =							torch.tensor(
           [[[9.2595, -3.6038, 11.8819], [9.3869, -3.2693, 11.0956], [11.8524, -3.4938, 13.1210]]] )

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






     @slow
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	List[str]      =							MraForMaskedLM.from_pretrained("""uw-madison/mra-base-4096-8-d3""" )
       a_      :	Dict      =							torch.arange(40_96 ).unsqueeze(0 )

       with torch.no_grad():
         a_      :	str      =							model(_lowerCAmelCase )[0]

       a_      :	str      =							5_02_65

       a_      :	List[str]      =							torch.Size((1, 40_96, vocab_size) )
       self.assertEqual(output.shape							,		_lowerCAmelCase )

       a_      :	Any      =							torch.tensor(
           [[[5.4789, -2.3564, 7.5064], [7.9067, -1.3369, 9.9668], [9.0712, -1.8106, 7.0380]]] )

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

'''simple docstring'''

import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()


class  _UpperCAmelCase ( lowerCAmelCase__	,unittest.TestCase    ):

     """simple docstring"""

     a_    	=					PriorTransformer
     a_    	=					"hidden_states"


     @property
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	Optional[int]      =							4
       a_      :	List[str]      =							8
       a_      :	List[str]      =							7

       a_      :	Tuple      =							floats_tensor((batch_size, embedding_dim) ).to(lowerCAmelCase_ )

       a_      :	List[Any]      =							floats_tensor((batch_size, embedding_dim) ).to(lowerCAmelCase_ )
       a_      :	List[Any]      =							floats_tensor((batch_size, num_embeddings, embedding_dim) ).to(lowerCAmelCase_ )

       return {
           "hidden_states": hidden_states,
           "timestep": 2,
           "proj_embedding": proj_embedding,
           "encoder_hidden_states": encoder_hidden_states,
       }


     def 							_lowerCAmelCase  (      self							,		lowerCAmelCase_=0 ):

       '''simple docstring'''

       torch.manual_seed(lowerCAmelCase_ )
       a_      :	List[Any]      =							4
       a_      :	Union[str, Any]      =							8
       a_      :	Optional[int]      =							7

       a_      :	List[Any]      =							torch.randn((batch_size, embedding_dim) ).to(lowerCAmelCase_ )

       a_      :	Optional[Any]      =							torch.randn((batch_size, embedding_dim) ).to(lowerCAmelCase_ )
       a_      :	Optional[Any]      =							torch.randn((batch_size, num_embeddings, embedding_dim) ).to(lowerCAmelCase_ )

       return {
           "hidden_states": hidden_states,
           "timestep": 2,
           "proj_embedding": proj_embedding,
           "encoder_hidden_states": encoder_hidden_states,
       }


     @property
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       return (4, 8)


     @property
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       return (4, 8)


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	Union[str, Any]      =							{
           """num_attention_heads""": 2,
           """attention_head_dim""": 4,
           """num_layers""": 2,
           """embedding_dim""": 8,
           """num_embeddings""": 7,
           """additional_embeddings""": 4,
       }
       a_      :	Optional[int]      =							self.dummy_input
       return init_dict, inputs_dict


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      , a_      :	Optional[Any]      =							PriorTransformer.from_pretrained(
           """hf-internal-testing/prior-dummy"""							,		output_loading_info=lowerCAmelCase_ )
       self.assertIsNotNone(lowerCAmelCase_ )
       self.assertEqual(len(loading_info["""missing_keys"""] )							,		0 )

       model.to(lowerCAmelCase_ )
       a_      :	List[Any]      =							model(**self.dummy_input )[0]

       assert hidden_states is not None, "Make sure output is not None"


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      , a_      :	Optional[int]      =							self.prepare_init_args_and_inputs_for_common()

       a_      :	str      =							self.model_class(**lowerCAmelCase_ )
       a_      :	Optional[Any]      =							inspect.signature(model.forward )
       # signature.parameters is an OrderedDict => so arg_names order is deterministic
       a_      :	str      =							[*signature.parameters.keys()]

       a_      :	str      =							["""hidden_states""", """timestep"""]
       self.assertListEqual(arg_names[:2]							,		lowerCAmelCase_ )






     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	Optional[Any]      =							PriorTransformer.from_pretrained("""hf-internal-testing/prior-dummy""" )
       a_      :	Optional[Any]      =							model.to(lowerCAmelCase_ )

       if hasattr(lowerCAmelCase_							,		"""set_default_attn_processor""" ):
         model.set_default_attn_processor()

       a_      :	Tuple      =							self.get_dummy_seed_input()

       with torch.no_grad():
         a_      :	List[str]      =							model(**lowerCAmelCase_ )[0]

       a_      :	Optional[Any]      =							output[0, :5].flatten().cpu()
       print(lowerCAmelCase_ )

       # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
       # the expected output slices are not the same for CPU and GPU.
       a_      :	List[str]      =							torch.tensor([-1.3436, -0.2870, 0.7538, 0.4368, -0.0239] )
       self.assertTrue(torch_all_close(lowerCAmelCase_							,		lowerCAmelCase_							,		rtol=1E-2 ) )







@slow
class  _UpperCAmelCase ( unittest.TestCase    ):

     """simple docstring"""

     def 							_lowerCAmelCase  (      self							,		lowerCAmelCase_=1							,		lowerCAmelCase_=7_68							,		lowerCAmelCase_=77							,		lowerCAmelCase_=0 ):

       '''simple docstring'''

       torch.manual_seed(lowerCAmelCase_ )
       a_      :	str      =							batch_size
       a_      :	Dict      =							embedding_dim
       a_      :	Tuple      =							num_embeddings

       a_      :	Optional[Any]      =							torch.randn((batch_size, embedding_dim) ).to(lowerCAmelCase_ )

       a_      :	List[str]      =							torch.randn((batch_size, embedding_dim) ).to(lowerCAmelCase_ )
       a_      :	List[Any]      =							torch.randn((batch_size, num_embeddings, embedding_dim) ).to(lowerCAmelCase_ )

       return {
           "hidden_states": hidden_states,
           "timestep": 2,
           "proj_embedding": proj_embedding,
           "encoder_hidden_states": encoder_hidden_states,
       }


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






     @parameterized.expand(
         [
             # fmt: off
             [13, [-0.5861,  0.1283, -0.0931,  0.0882,  0.4476,  0.1329, -0.0498,  0.0640]],
             [37, [-0.4913,  0.0110, -0.0483,  0.0541,  0.4954, -0.0170,  0.0354,  0.1651]],
             # fmt: on
         ] )
     def 							_lowerCAmelCase  (      self							,		lowerCAmelCase_							,		lowerCAmelCase_ ):

       '''simple docstring'''

       a_      :	Any      =							PriorTransformer.from_pretrained("""kandinsky-community/kandinsky-2-1-prior"""							,		subfolder="""prior""" )
       model.to(lowerCAmelCase_ )
       a_      :	Optional[int]      =							self.get_dummy_seed_input(seed=lowerCAmelCase_ )

       with torch.no_grad():
         a_      :	str      =							model(**lowerCAmelCase_ )[0]

       assert list(sample.shape ) == [1, 7_68]

       a_      :	int      =							sample[0, :8].flatten().cpu()
       print(lowerCAmelCase_ )
       a_      :	str      =							torch.tensor(lowerCAmelCase_ )

       assert torch_all_close(lowerCAmelCase_							,		lowerCAmelCase_							,		atol=1E-3 )


