

import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

import diffusers
from diffusers import (
    AutoencoderKL,
    EulerDiscreteScheduler,
    StableDiffusionLatentUpscalePipeline,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


def 					_UpperCAmelCase				(						a__):

			'''simple docstring'''



			a_						:						int 							= [tensor.shape for tensor in tensor_list]
			return all(shape == shapes[0] for shape in shapes[1:])



class        A__(a_,  a_,  a_,  unittest.TestCase							):
						"""simple docstring"""






						_A    :				Dict       					=			StableDiffusionLatentUpscalePipeline
						_A    :				Dict       					=			TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
						    '''height''',
						    '''width''',
						    '''cross_attention_kwargs''',
						    '''negative_prompt_embeds''',
						    '''prompt_embeds''',
						}
						_A    :				str       					=			PipelineTesterMixin.required_optional_params - {'''num_images_per_prompt'''}
						_A    :				Any       					=			TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
						_A    :				int       					=			frozenset(
						    []							)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
						_A    :				Union[str, Any]       					=			frozenset([]							)

						_A    :				Optional[int]       					=			True





						@property
						def 			UpperCamelCase__					(    self    )     -> Union[str, Any]:
									a_						:						Optional[int] 							= 1
									a_						:						Any 							= 4
									a_						:						int 							= (16, 16)

									a_						:						List[str] 							= floats_tensor((batch_size, num_channels) + sizes  ,      rng=random.Random(0    )    ).to(_lowercase    )
									return image





						def 			UpperCamelCase__					(    self    )     -> List[Any]:
									torch.manual_seed(0    )
									a_						:						Optional[Any] 							= UNetaDConditionModel(
									    act_fn="""gelu"""  ,      attention_head_dim=8  ,      norm_num_groups=_lowercase  ,      block_out_channels=[32, 32, 64, 64]  ,      time_cond_proj_dim=160  ,      conv_in_kernel=1  ,      conv_out_kernel=1  ,      cross_attention_dim=32  ,      down_block_types=(
									        """KDownBlock2D""",
									        """KCrossAttnDownBlock2D""",
									        """KCrossAttnDownBlock2D""",
									        """KCrossAttnDownBlock2D""",
									    )  ,      in_channels=8  ,      mid_block_type=_lowercase  ,      only_cross_attention=_lowercase  ,      out_channels=5  ,      resnet_time_scale_shift="""scale_shift"""  ,      time_embedding_type="""fourier"""  ,      timestep_post_act="""gelu"""  ,      up_block_types=("""KCrossAttnUpBlock2D""", """KCrossAttnUpBlock2D""", """KCrossAttnUpBlock2D""", """KUpBlock2D""")  ,      )
									a_						:						Optional[Any] 							= AutoencoderKL(
									    block_out_channels=[32, 32, 64, 64]  ,      in_channels=3  ,      out_channels=3  ,      down_block_types=[
									        """DownEncoderBlock2D""",
									        """DownEncoderBlock2D""",
									        """DownEncoderBlock2D""",
									        """DownEncoderBlock2D""",
									    ]  ,      up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D""", """UpDecoderBlock2D""", """UpDecoderBlock2D"""]  ,      latent_channels=4  ,      )
									a_						:						Any 							= EulerDiscreteScheduler(prediction_type="""sample"""    )
									a_						:						int 							= CLIPTextConfig(
									    bos_token_id=0  ,      eos_token_id=2  ,      hidden_size=32  ,      intermediate_size=37  ,      layer_norm_eps=1e-05  ,      num_attention_heads=4  ,      num_hidden_layers=5  ,      pad_token_id=1  ,      vocab_size=1_000  ,      hidden_act="""quick_gelu"""  ,      projection_dim=512  ,      )
									a_						:						str 							= CLIPTextModel(_lowercase    )
									a_						:						Optional[int] 							= CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""    )

									a_						:						Tuple 							= {
									    """unet""": model.eval(),
									    """vae""": vae.eval(),
									    """scheduler""": scheduler,
									    """text_encoder""": text_encoder,
									    """tokenizer""": tokenizer,
									}

									return components





						def 			UpperCamelCase__					(    self  ,      _lowercase  ,      _lowercase=0    )     -> List[Any]:
									if str(_lowercase    ).startswith("""mps"""    ):
												a_						:						Tuple 							= torch.manual_seed(_lowercase    )
									else:
												a_						:						Optional[int] 							= torch.Generator(device=_lowercase    ).manual_seed(_lowercase    )
									a_						:						Tuple 							= {
									    """prompt""": """A painting of a squirrel eating a burger""",
									    """image""": self.dummy_image.cpu(),
									    """generator""": generator,
									    """num_inference_steps""": 2,
									    """output_type""": """numpy""",
									}
									return inputs





						def 			UpperCamelCase__					(    self    )     -> int:
									a_						:						int 							= """cpu"""

									a_						:						List[str] 							= self.get_dummy_components()
									a_						:						str 							= self.pipeline_class(**_lowercase    )
									pipe.to(_lowercase    )
									pipe.set_progress_bar_config(disable=_lowercase    )

									a_						:						List[Any] 							= self.get_dummy_inputs(_lowercase    )
									a_						:						List[Any] 							= pipe(**_lowercase    ).images
									a_						:						Any 							= image[0, -3:, -3:, -1]

									self.assertEqual(image.shape  ,      (1, 256, 256, 3)    )
									a_						:						Dict 							= np.array(
									    [0.4_7_2_2_2_4_1_2, 0.4_1_9_2_1_6_3_3, 0.4_4_7_1_7_4_3_4, 0.4_6_8_7_4_1_9_2, 0.4_2_5_8_8_2_5_8, 0.4_6_1_5_0_7_2_6, 0.4_6_7_7_5_3_4, 0.4_5_5_8_3_8_3_2, 0.4_8_5_7_9_0_5_5]    )
									a_						:						Dict 							= np.abs(image_slice.flatten() - expected_slice    ).max()
									self.assertLessEqual(_lowercase  ,      1e-3    )





						def 			UpperCamelCase__					(    self    )     -> Tuple:
									super().test_attention_slicing_forward_pass(expected_max_diff=7e-3    )





						def 			UpperCamelCase__					(    self    )     -> Tuple:
									super().test_cpu_offload_forward_pass(expected_max_diff=3e-3    )





						def 			UpperCamelCase__					(    self    )     -> List[str]:
									super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3    )





						def 			UpperCamelCase__					(    self    )     -> Dict:
									super().test_inference_batch_single_identical(expected_max_diff=7e-3    )





						def 			UpperCamelCase__					(    self    )     -> Tuple:
									super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3e-3    )





						def 			UpperCamelCase__					(    self    )     -> Union[str, Any]:
									super().test_save_load_local(expected_max_difference=3e-3    )





						def 			UpperCamelCase__					(    self    )     -> str:
									super().test_save_load_optional_components(expected_max_difference=3e-3    )





						def 			UpperCamelCase__					(    self    )     -> Any:
									a_						:						str 							= [
									    """DDIMScheduler""",
									    """DDPMScheduler""",
									    """PNDMScheduler""",
									    """HeunDiscreteScheduler""",
									    """EulerAncestralDiscreteScheduler""",
									    """KDPM2DiscreteScheduler""",
									    """KDPM2AncestralDiscreteScheduler""",
									    """DPMSolverSDEScheduler""",
									]
									a_						:						Union[str, Any] 							= self.get_dummy_components()
									a_						:						int 							= self.pipeline_class(**_lowercase    )

									# make sure that PNDM does not need warm-up
									pipe.scheduler.register_to_config(skip_prk_steps=_lowercase    )

									pipe.to(_lowercase    )
									pipe.set_progress_bar_config(disable=_lowercase    )
									a_						:						Any 							= self.get_dummy_inputs(_lowercase    )
									a_						:						List[str] 							= 2

									a_						:						Optional[Any] 							= []
									for scheduler_enum in KarrasDiffusionSchedulers:
												if scheduler_enum.name in skip_schedulers:
															# no sigma schedulers are not supported
															# no schedulers
															continue

												a_						:						Any 							= getattr(_lowercase  ,      scheduler_enum.name    )
												a_						:						Optional[Any] 							= scheduler_cls.from_config(pipe.scheduler.config    )
												a_						:						Optional[int] 							= pipe(**_lowercase    )[0]
												outputs.append(_lowercase    )

									assert check_same_shape(_lowercase    )







@require_torch_gpu
@slow
class        A__(unittest.TestCase							):
						"""simple docstring"""






						def 			UpperCamelCase__					(    self    )     -> List[Any]:
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()





						def 			UpperCamelCase__					(    self    )     -> Any:
									a_						:						Dict 							= torch.manual_seed(33    )

									a_						:						str 							= StableDiffusionPipeline.from_pretrained("""CompVis/stable-diffusion-v1-4"""  ,      torch_dtype=torch.floataa    )
									pipe.to("""cuda"""    )

									a_						:						int 							= StableDiffusionLatentUpscalePipeline.from_pretrained(
									    """stabilityai/sd-x2-latent-upscaler"""  ,      torch_dtype=torch.floataa    )
									upscaler.to("""cuda"""    )

									a_						:						Optional[int] 							= """a photo of an astronaut high resolution, unreal engine, ultra realistic"""

									a_						:						Tuple 							= pipe(_lowercase  ,      generator=_lowercase  ,      output_type="""latent"""    ).images

									a_						:						Tuple 							= upscaler(
									    prompt=_lowercase  ,      image=_lowercase  ,      num_inference_steps=20  ,      guidance_scale=0  ,      generator=_lowercase  ,      output_type="""np"""  ,      ).images[0]

									a_						:						List[Any] 							= load_numpy(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/astronaut_1024.npy"""    )
									assert np.abs((expected_image - image).mean()    ) < 5e-2





						def 			UpperCamelCase__					(    self    )     -> Dict:
									a_						:						Optional[int] 							= torch.manual_seed(33    )

									a_						:						Tuple 							= StableDiffusionLatentUpscalePipeline.from_pretrained(
									    """stabilityai/sd-x2-latent-upscaler"""  ,      torch_dtype=torch.floataa    )
									upscaler.to("""cuda"""    )

									a_						:						List[str] 							= """the temple of fire by Ross Tran and Gerardo Dottori, oil on canvas"""

									a_						:						Optional[Any] 							= load_image(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_512.png"""    )

									a_						:						Any 							= upscaler(
									    prompt=_lowercase  ,      image=_lowercase  ,      num_inference_steps=20  ,      guidance_scale=0  ,      generator=_lowercase  ,      output_type="""np"""  ,      ).images[0]

									a_						:						Optional[Any] 							= load_numpy(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_1024.npy"""    )
									assert np.abs((expected_image - image).max()    ) < 5e-2





import unittest

from parameterized import parameterized

from transformers import OpenLlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import OpenLlamaForCausalLM, OpenLlamaForSequenceClassification, OpenLlamaModel



class        A__:
						"""simple docstring"""






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





						def 			UpperCamelCase__					(    self    )     -> Any:
									a_						:						Dict 							= ids_tensor([self.batch_size, self.seq_length]  ,      self.vocab_size    )

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

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

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

									a_						:						Optional[int] 							= self.get_config()

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





						def 			UpperCamelCase__					(    self    )     -> Any:
									return OpenLlamaConfig(
									    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  ,      use_stable_embedding=_lowercase  ,      )





						def 			UpperCamelCase__					(    self  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase    )     -> Dict:
									a_						:						Optional[Any] 							= OpenLlamaModel(config=_lowercase    )
									model.to(_lowercase    )
									model.eval()
									a_						:						List[Any] 							= model(_lowercase  ,      attention_mask=_lowercase    )
									a_						:						Optional[Any] 							= 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  ,      )     -> Union[str, Any]:
									a_						:						int 							= True
									a_						:						Union[str, Any] 							= OpenLlamaModel(_lowercase    )
									model.to(_lowercase    )
									model.eval()
									a_						:						Optional[int] 							= model(
									    _lowercase  ,      attention_mask=_lowercase  ,      encoder_hidden_states=_lowercase  ,      encoder_attention_mask=_lowercase  ,      )
									a_						:						Any 							= model(
									    _lowercase  ,      attention_mask=_lowercase  ,      encoder_hidden_states=_lowercase  ,      )
									a_						:						Optional[int] 							= model(_lowercase  ,      attention_mask=_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  ,      )     -> Any:
									a_						:						Dict 							= OpenLlamaForCausalLM(config=_lowercase    )
									model.to(_lowercase    )
									model.eval()
									a_						:						int 							= model(_lowercase  ,      attention_mask=_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  ,      _lowercase  ,      _lowercase  ,      )     -> List[str]:
									a_						:						Dict 							= True
									a_						:						Optional[int] 							= True
									a_						:						Dict 							= OpenLlamaForCausalLM(config=_lowercase    )
									model.to(_lowercase    )
									model.eval()

									# first forward pass
									a_						:						List[str] 							= model(
									    _lowercase  ,      attention_mask=_lowercase  ,      encoder_hidden_states=_lowercase  ,      encoder_attention_mask=_lowercase  ,      use_cache=_lowercase  ,      )
									a_						:						Optional[Any] 							= outputs.past_key_values

									# create hypothetical multiple next token and extent to next_input_ids
									a_						:						Optional[Any] 							= ids_tensor((self.batch_size, 3)  ,      config.vocab_size    )
									a_						:						Union[str, Any] 							= ids_tensor((self.batch_size, 3)  ,      vocab_size=2    )

									# append to next input_ids and
									a_						:						Any 							= torch.cat([input_ids, next_tokens]  ,      dim=-1    )
									a_						:						Any 							= torch.cat([input_mask, next_mask]  ,      dim=-1    )

									a_						:						Optional[int] 							= model(
									    _lowercase  ,      attention_mask=_lowercase  ,      encoder_hidden_states=_lowercase  ,      encoder_attention_mask=_lowercase  ,      output_hidden_states=_lowercase  ,      )["""hidden_states"""][0]
									a_						:						Dict 							= model(
									    _lowercase  ,      attention_mask=_lowercase  ,      encoder_hidden_states=_lowercase  ,      encoder_attention_mask=_lowercase  ,      past_key_values=_lowercase  ,      output_hidden_states=_lowercase  ,      )["""hidden_states"""][0]

									# select random slice
									a_						:						Optional[int] 							= ids_tensor((1,)  ,      output_from_past.shape[-1]    ).item()
									a_						:						Dict 							= output_from_no_past[:, -3:, random_slice_idx].detach()
									a_						:						Optional[int] 							= output_from_past[:, :, random_slice_idx].detach()

									self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]    )

									# test that outputs are equal for slice
									self.parent.assertTrue(torch.allclose(_lowercase  ,      _lowercase  ,      atol=1e-3    )    )





						def 			UpperCamelCase__					(    self    )     -> Optional[Any]:
									a_						:						Any 							= self.prepare_config_and_inputs()
									(
									    (
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					
									)						:						Union[str, Any] 							= config_and_inputs
									a_						:						Tuple 							= {"""input_ids""": input_ids, """attention_mask""": input_mask}
									return config, inputs_dict







@require_torch
class        A__(a_,  a_,  a_,  unittest.TestCase							):
						"""simple docstring"""






						_A    :				Tuple       					=			(
						    (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
						)
						_A    :				int       					=			(OpenLlamaForCausalLM,) if is_torch_available() else ()
						_A    :				Optional[int]       					=			(
						    {
						        '''feature-extraction''': OpenLlamaModel,
						        '''text-classification''': OpenLlamaForSequenceClassification,
						        '''text-generation''': OpenLlamaForCausalLM,
						        '''zero-shot''': OpenLlamaForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)
						_A    :				Tuple       					=			False
						_A    :				Any       					=			False





						def 			UpperCamelCase__					(    self    )     -> List[str]:
									a_						:						Optional[int] 							= OpenLlamaModelTester(self    )
									a_						:						Optional[int] 							= ConfigTester(self  ,      config_class=_lowercase  ,      hidden_size=37    )





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





						def 			UpperCamelCase__					(    self    )     -> Optional[Any]:
									a_						:						Union[str, Any] 							= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*_lowercase    )





						def 			UpperCamelCase__					(    self    )     -> Tuple:
									a_						:						Any 							= self.model_tester.prepare_config_and_inputs()
									for type in ["absolute", "relative_key", "relative_key_query"]:
												a_						:						Dict 							= type
												self.model_tester.create_and_check_model(*_lowercase    )





						def 			UpperCamelCase__					(    self    )     -> Union[str, Any]:
									a_							,					a_						:						Tuple 							= self.model_tester.prepare_config_and_inputs_for_common()
									a_						:						Any 							= 3
									a_						:						List[str] 							= input_dict["""input_ids"""]
									a_						:						List[str] 							= input_ids.ne(1    ).to(_lowercase    )
									a_						:						int 							= ids_tensor([self.model_tester.batch_size]  ,      self.model_tester.type_sequence_label_size    )
									a_						:						int 							= OpenLlamaForSequenceClassification(_lowercase    )
									model.to(_lowercase    )
									model.eval()
									a_						:						Tuple 							= model(_lowercase  ,      attention_mask=_lowercase  ,      labels=_lowercase    )
									self.assertEqual(result.logits.shape  ,      (self.model_tester.batch_size, self.model_tester.num_labels)    )





						def 			UpperCamelCase__					(    self    )     -> Optional[Any]:
									a_							,					a_						:						Any 							= self.model_tester.prepare_config_and_inputs_for_common()
									a_						:						Union[str, Any] 							= 3
									a_						:						List[Any] 							= """single_label_classification"""
									a_						:						Dict 							= input_dict["""input_ids"""]
									a_						:						int 							= input_ids.ne(1    ).to(_lowercase    )
									a_						:						List[str] 							= ids_tensor([self.model_tester.batch_size]  ,      self.model_tester.type_sequence_label_size    )
									a_						:						Any 							= OpenLlamaForSequenceClassification(_lowercase    )
									model.to(_lowercase    )
									model.eval()
									a_						:						Any 							= model(_lowercase  ,      attention_mask=_lowercase  ,      labels=_lowercase    )
									self.assertEqual(result.logits.shape  ,      (self.model_tester.batch_size, self.model_tester.num_labels)    )





						def 			UpperCamelCase__					(    self    )     -> Tuple:
									a_							,					a_						:						Union[str, Any] 							= self.model_tester.prepare_config_and_inputs_for_common()
									a_						:						Any 							= 3
									a_						:						List[str] 							= """multi_label_classification"""
									a_						:						Dict 							= input_dict["""input_ids"""]
									a_						:						Any 							= input_ids.ne(1    ).to(_lowercase    )
									a_						:						List[Any] 							= ids_tensor(
									    [self.model_tester.batch_size, config.num_labels]  ,      self.model_tester.type_sequence_label_size    ).to(torch.float    )
									a_						:						Optional[int] 							= OpenLlamaForSequenceClassification(_lowercase    )
									model.to(_lowercase    )
									model.eval()
									a_						:						int 							= model(_lowercase  ,      attention_mask=_lowercase  ,      labels=_lowercase    )
									self.assertEqual(result.logits.shape  ,      (self.model_tester.batch_size, self.model_tester.num_labels)    )





						@unittest.skip("""Open-Llama buffers include complex numbers, which breaks this test"""    )
						def 			UpperCamelCase__					(    self    )     -> Tuple:
									pass





						@parameterized.expand([("""linear""",), ("""dynamic""",)]    )
						def 			UpperCamelCase__					(    self  ,      _lowercase    )     -> str:
									a_							,					a_						:						Optional[int] 							= self.model_tester.prepare_config_and_inputs_for_common()
									a_						:						List[Any] 							= ids_tensor([1, 10]  ,      config.vocab_size    )
									a_						:						str 							= ids_tensor([1, int(config.max_position_embeddings * 1.5    )]  ,      config.vocab_size    )

									set_seed(42    )  # Fixed seed at init time so the two models get the same random weights
									a_						:						Union[str, Any] 							= OpenLlamaModel(_lowercase    )
									original_model.to(_lowercase    )
									original_model.eval()
									a_						:						Union[str, Any] 							= original_model(_lowercase    ).last_hidden_state
									a_						:						str 							= original_model(_lowercase    ).last_hidden_state

									set_seed(42    )  # Fixed seed at init time so the two models get the same random weights
									a_						:						int 							= {"""type""": scaling_type, """factor""": 1_0.0}
									a_						:						Union[str, Any] 							= OpenLlamaModel(_lowercase    )
									scaled_model.to(_lowercase    )
									scaled_model.eval()
									a_						:						Optional[int] 							= scaled_model(_lowercase    ).last_hidden_state
									a_						:						Tuple 							= scaled_model(_lowercase    ).last_hidden_state

									# Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
									# maximum sequence length, so the outputs for the short input should match.
									if scaling_type == "dynamic":
												self.assertTrue(torch.allclose(_lowercase  ,      _lowercase  ,      atol=1e-5    )    )
									else:
												self.assertFalse(torch.allclose(_lowercase  ,      _lowercase  ,      atol=1e-5    )    )

									# The output should be different for long inputs
									self.assertFalse(torch.allclose(_lowercase  ,      _lowercase  ,      atol=1e-5    )    )


