


from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


snake_case				:    Any       				=     logging.get_logger(__name__)  # pylint: disable=invalid-name

snake_case				:    Union[str, Any]       				=     '''
    Examples:
        ```py
        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior")
        >>> pipe_prior.to("cuda")
        >>> prompt = "red cat, 4k photo"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> zero_image_emb = out.negative_image_embeds
        >>> pipe = KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder")
        >>> pipe.to("cuda")
        >>> image = pipe(
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=50,
        ... ).images
        >>> image[0].save("cat.png")
        ```
'''
def 				__lowerCamelCase						( UpperCAmelCase_						:		List[Any] ,  UpperCAmelCase_						:		Optional[int] ,  UpperCAmelCase_						:		List[str]=8						):




      """simple docstring"""





      a						:List[str] 							=       height // scale_factor**2
      if height % scale_factor**2 != 0:
            new_height += 1
      a						:int 							=       width // scale_factor**2
      if width % scale_factor**2 != 0:
            new_width += 1
      return new_height * scale_factor, new_width * scale_factor







class 		_snake_case  (      _snake_case							):




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

       self.register_modules(
           unet=_lowerCamelCase					,  scheduler=_lowerCamelCase					,  movq=_lowerCamelCase					,  )
       a						:Any 							=       2 ** (len(self.movq.config.block_out_channels    ) - 1)




 def        SCREAMING_SNAKE_CASE__				(		self					,  _lowerCamelCase					,  _lowerCamelCase					,  _lowerCamelCase					,  _lowerCamelCase					,  _lowerCamelCase					,  _lowerCamelCase    ):
       if latents is None:
             a						:str 							=       randn_tensor(_lowerCamelCase					,  generator=_lowerCamelCase					,  device=_lowerCamelCase					,  dtype=_lowerCamelCase    )
       else:
             if latents.shape != shape:
                   raise ValueError(F'''Unexpected latents shape, got {latents.shape}, expected {shape}'''    )
             a						:Any 							=       latents.to(_lowerCamelCase    )

       a						:Dict 							=       latents * scheduler.init_noise_sigma
       return latents




 def        SCREAMING_SNAKE_CASE__				(		self					,  _lowerCamelCase=0    ):
       if is_accelerate_available():
             from accelerate import cpu_offload
       else:
             raise ImportError('''Please install accelerate via `pip install accelerate`'''    )

       a						:int 							=       torch.device(F'''cuda:{gpu_id}'''    )

       a						:int 							=       [
           self.unet,
           self.movq,
       ]
       for cpu_offloaded_model in models:
             if cpu_offloaded_model is not None:
                   cpu_offload(_lowerCamelCase					,  _lowerCamelCase    )




 def        SCREAMING_SNAKE_CASE__				(		self					,  _lowerCamelCase=0    ):
       if is_accelerate_available() and is_accelerate_version('''>='''					,  '''0.17.0.dev0'''    ):
             from accelerate import cpu_offload_with_hook
       else:
             raise ImportError('''`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'''    )

       a						:Any 							=       torch.device(F'''cuda:{gpu_id}'''    )

       if self.device.type != "cpu":
             self.to('''cpu'''					,  silence_dtype_warnings=_lowerCamelCase    )
             torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

       a						:Tuple 							=       None
       for cpu_offloaded_model in [self.unet, self.movq]:
             a							,				a						:List[str] 							=       cpu_offload_with_hook(_lowerCamelCase					,  _lowerCamelCase					,  prev_module_hook=_lowerCamelCase    )

       # We'll offload the last model manually.
       a						:str 							=       hook




 @property
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
 def        SCREAMING_SNAKE_CASE__				(		self    ):
       if not hasattr(self.unet					,  '''_hf_hook'''    ):
             return self.device
       for module in self.unet.modules():
             if (
                 hasattr(_lowerCamelCase					,  '''_hf_hook'''    )
                 and hasattr(module._hf_hook					,  '''execution_device'''    )
                 and module._hf_hook.execution_device is not None
             ):
                   return torch.device(module._hf_hook.execution_device    )
       return self.device




 @torch.no_grad()
 @replace_example_docstring(_lowerCamelCase    )
 def __call__(		self					,  _lowerCamelCase					,  _lowerCamelCase					,  _lowerCamelCase = 512					,  _lowerCamelCase = 512					,  _lowerCamelCase = 100					,  _lowerCamelCase = 4.0					,  _lowerCamelCase = 1					,  _lowerCamelCase = None					,  _lowerCamelCase = None					,  _lowerCamelCase = "pil"					,  _lowerCamelCase = True					,  ):
       a						:int 							=       self._execution_device

       a						:Optional[Any] 							=       guidance_scale > 1.0

       if isinstance(_lowerCamelCase					,  _lowerCamelCase    ):
             a						:Union[str, Any] 							=       torch.cat(_lowerCamelCase					,  dim=0    )
       a						:Any 							=       image_embeds.shape[0] * num_images_per_prompt
       if isinstance(_lowerCamelCase					,  _lowerCamelCase    ):
             a						:List[str] 							=       torch.cat(_lowerCamelCase					,  dim=0    )

       if do_classifier_free_guidance:
             a						:Union[str, Any] 							=       image_embeds.repeat_interleave(_lowerCamelCase					,  dim=0    )
             a						:Optional[int] 							=       negative_image_embeds.repeat_interleave(_lowerCamelCase					,  dim=0    )

       a						:Optional[int] 							=       torch.cat([negative_image_embeds, image_embeds]					,  dim=0    ).to(dtype=self.unet.dtype					,  device=_lowerCamelCase    )

       self.scheduler.set_timesteps(_lowerCamelCase					,  device=_lowerCamelCase    )
       a						:Optional[Any] 							=       self.scheduler.timesteps

       a						:List[str] 							=       self.unet.config.in_channels

       a							,				a						:str 							=       downscale_height_and_width(_lowerCamelCase					,  _lowerCamelCase					,  self.movq_scale_factor    )

       # create initial latent
       a						:int 							=       self.prepare_latents(
           (batch_size, num_channels_latents, height, width)					,  image_embeds.dtype					,  _lowerCamelCase					,  _lowerCamelCase					,  _lowerCamelCase					,  self.scheduler					,  )

       for i, t in enumerate(self.progress_bar(_lowerCamelCase    )    ):
             # expand the latents if we are doing classifier free guidance
             a						:Tuple 							=       torch.cat([latents] * 2    ) if do_classifier_free_guidance else latents

             a						:Union[str, Any] 							=       {'''image_embeds''': image_embeds}
             a						:Optional[Any] 							=       self.unet(
                 sample=_lowerCamelCase					,  timestep=_lowerCamelCase					,  encoder_hidden_states=_lowerCamelCase					,  added_cond_kwargs=_lowerCamelCase					,  return_dict=_lowerCamelCase					,  )[0]

             if do_classifier_free_guidance:
                   a							,				a						:Any 							=       noise_pred.split(latents.shape[1]					,  dim=1    )
                   a							,				a						:List[str] 							=       noise_pred.chunk(2    )
                   a							,				a						:int 							=       variance_pred.chunk(2    )
                   a						:List[Any] 							=       noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                   a						:Optional[int] 							=       torch.cat([noise_pred, variance_pred_text]					,  dim=1    )

             if not (
                 hasattr(self.scheduler.config					,  '''variance_type'''    )
                 and self.scheduler.config.variance_type in ["learned", "learned_range"]
             ):
                   a							,				a						:Tuple 							=       noise_pred.split(latents.shape[1]					,  dim=1    )

             # compute the previous noisy sample x_t -> x_t-1
             a						:int 							=       self.scheduler.step(
                 _lowerCamelCase					,  _lowerCamelCase					,  _lowerCamelCase					,  generator=_lowerCamelCase					,  )[0]
       # post-processing
       a						:int 							=       self.movq.decode(_lowerCamelCase					,  force_not_quantize=_lowerCamelCase    )['''sample''']

       if output_type not in ["pt", "np", "pil"]:
             raise ValueError(F'''Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}'''    )

       if output_type in ["np", "pil"]:
             a						:str 							=       image * 0.5 + 0.5
             a						:List[Any] 							=       image.clamp(0					,  1    )
             a						:str 							=       image.cpu().permute(0					,  2					,  3					,  1    ).float().numpy()

       if output_type == "pil":
             a						:str 							=       self.numpy_to_pil(_lowerCamelCase    )

       if not return_dict:
             return (image,)

       return ImagePipelineOutput(images=_lowerCamelCase    )








"""simple docstring"""


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

				def __init__(	self							:				Any     ,  UpperCamelCase__							:				int     ,  UpperCamelCase__							:				str=1_3     ,  UpperCamelCase__							:				Tuple=7     ,  UpperCamelCase__							:				Optional[Any]=True     ,  UpperCamelCase__							:				Optional[Any]=True     ,  UpperCamelCase__							:				Any=False     ,  UpperCamelCase__							:				str=True     ,  UpperCamelCase__							:				List[Any]=9_9     ,  UpperCamelCase__							:				Optional[int]=3_2     ,  UpperCamelCase__							:				List[Any]=5     ,  UpperCamelCase__							:				Union[str, Any]=4     ,  UpperCamelCase__							:				List[str]=3_7     ,  UpperCamelCase__							:				Optional[Any]="gelu"     ,  UpperCamelCase__							:				int=0.1     ,  UpperCamelCase__							:				Dict=0.1     ,  UpperCamelCase__							:				Dict=5_1_2     ,  UpperCamelCase__							:				List[Any]=1_6     ,  UpperCamelCase__							:				Optional[int]=2     ,  UpperCamelCase__							:				Dict=0.02     ,  UpperCamelCase__							:				int=3     ,  UpperCamelCase__							:				Any=4     ,  UpperCamelCase__							:				Optional[int]=None     ,  )->   Optional[Any]:



						'''simple docstring'''


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

				def 							lowercase_     (	self							:				Optional[Any])->   Optional[Any]:



						'''simple docstring'''


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

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

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

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

						__lowerCAmelCase:       Any          =			self.get_config()

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

				def 							lowercase_     (	self							:				str)->   Dict:



						'''simple docstring'''


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

				def 							lowercase_     (	self							:				Optional[Any]     ,  UpperCamelCase__							:				int     ,  UpperCamelCase__							:				Tuple     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				Tuple     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				Optional[Any])->   str:



						'''simple docstring'''


						__lowerCAmelCase:       List[Any]          =			OpenLlamaModel(config=UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase:       Union[str, Any]          =			model(UpperCamelCase__     ,  attention_mask=UpperCamelCase__)
						__lowerCAmelCase:       int          =			model(UpperCamelCase__)
						self.parent.assertEqual(result.last_hidden_state.shape     ,  (self.batch_size, self.seq_length, self.hidden_size))

				def 							lowercase_     (	self							:				Optional[Any]     ,  UpperCamelCase__							:				List[str]     ,  UpperCamelCase__							:				Dict     ,  UpperCamelCase__							:				Union[str, Any]     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				Any     ,  UpperCamelCase__							:				Any     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				Union[str, Any]     ,  UpperCamelCase__							:				Tuple     ,  )->   Optional[int]:



						'''simple docstring'''


						__lowerCAmelCase:       Optional[Any]          =			True
						__lowerCAmelCase:       Any          =			OpenLlamaModel(UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase:       List[Any]          =			model(
						    UpperCamelCase__     ,  attention_mask=UpperCamelCase__     ,  encoder_hidden_states=UpperCamelCase__     ,  encoder_attention_mask=UpperCamelCase__     ,  )
						__lowerCAmelCase:       Optional[int]          =			model(
						    UpperCamelCase__     ,  attention_mask=UpperCamelCase__     ,  encoder_hidden_states=UpperCamelCase__     ,  )
						__lowerCAmelCase:       Optional[Any]          =			model(UpperCamelCase__     ,  attention_mask=UpperCamelCase__)
						self.parent.assertEqual(result.last_hidden_state.shape     ,  (self.batch_size, self.seq_length, self.hidden_size))

				def 							lowercase_     (	self							:				Dict     ,  UpperCamelCase__							:				Optional[int]     ,  UpperCamelCase__							:				Optional[int]     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				Tuple     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				Union[str, Any]     ,  UpperCamelCase__							:				Union[str, Any]     ,  UpperCamelCase__							:				List[str]     ,  )->   Optional[int]:



						'''simple docstring'''


						__lowerCAmelCase:       str          =			OpenLlamaForCausalLM(config=UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase:       Optional[int]          =			model(UpperCamelCase__     ,  attention_mask=UpperCamelCase__     ,  labels=UpperCamelCase__)
						self.parent.assertEqual(result.logits.shape     ,  (self.batch_size, self.seq_length, self.vocab_size))

				def 							lowercase_     (	self							:				Tuple     ,  UpperCamelCase__							:				Union[str, Any]     ,  UpperCamelCase__							:				str     ,  UpperCamelCase__							:				Optional[int]     ,  UpperCamelCase__							:				Any     ,  UpperCamelCase__							:				int     ,  UpperCamelCase__							:				int     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				List[Any]     ,  UpperCamelCase__							:				int     ,  )->   Tuple:



						'''simple docstring'''


						__lowerCAmelCase:       Union[str, Any]          =			True
						__lowerCAmelCase:       Dict          =			True
						__lowerCAmelCase:       Union[str, Any]          =			OpenLlamaForCausalLM(config=UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()

						# first forward pass
						__lowerCAmelCase:       Optional[int]          =			model(
						    UpperCamelCase__     ,  attention_mask=UpperCamelCase__     ,  encoder_hidden_states=UpperCamelCase__     ,  encoder_attention_mask=UpperCamelCase__     ,  use_cache=UpperCamelCase__     ,  )
						__lowerCAmelCase:       Optional[int]          =			outputs.past_key_values

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

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

						__lowerCAmelCase:       Union[str, Any]          =			model(
						    UpperCamelCase__     ,  attention_mask=UpperCamelCase__     ,  encoder_hidden_states=UpperCamelCase__     ,  encoder_attention_mask=UpperCamelCase__     ,  output_hidden_states=UpperCamelCase__     ,  )["hidden_states"][0]
						__lowerCAmelCase:       List[str]          =			model(
						    UpperCamelCase__     ,  attention_mask=UpperCamelCase__     ,  encoder_hidden_states=UpperCamelCase__     ,  encoder_attention_mask=UpperCamelCase__     ,  past_key_values=UpperCamelCase__     ,  output_hidden_states=UpperCamelCase__     ,  )["hidden_states"][0]

						# select random slice
						__lowerCAmelCase:       List[str]          =			ids_tensor((1,)     ,  output_from_past.shape[-1]).item()
						__lowerCAmelCase:       List[Any]          =			output_from_no_past[:, -3:, random_slice_idx].detach()
						__lowerCAmelCase:       Union[str, Any]          =			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(UpperCamelCase__     ,  UpperCamelCase__     ,  atol=1e-3))




				def 							lowercase_     (	self							:				Tuple)->   str:



						'''simple docstring'''


						__lowerCAmelCase:       Union[str, Any]          =			self.prepare_config_and_inputs()
						(
						    (
						    __lowerCAmelCase
						)      , (
						    __lowerCAmelCase
						)      , (
						    __lowerCAmelCase
						)      , (
						    __lowerCAmelCase
						)      , (
						    __lowerCAmelCase
						)      , (
						    __lowerCAmelCase
						)      , (
						    __lowerCAmelCase
						)      , 
						):       List[Any]          =			config_and_inputs
						__lowerCAmelCase:       Any          =			{"input_ids": input_ids, "attention_mask": input_mask}
						return config, inputs_dict





@require_torch
class     snake_case	(				__snake_case,							__snake_case,							__snake_case,							unittest.TestCase  ):
				SCREAMING_SNAKE_CASE_   :      Tuple         =  (
				    (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
				)
				SCREAMING_SNAKE_CASE_   :      Any         =  (OpenLlamaForCausalLM,) if is_torch_available() else ()
				SCREAMING_SNAKE_CASE_   :      Optional[int]         =  (
				    {
				        """feature-extraction""": OpenLlamaModel,
				        """text-classification""": OpenLlamaForSequenceClassification,
				        """text-generation""": OpenLlamaForCausalLM,
				        """zero-shot""": OpenLlamaForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				SCREAMING_SNAKE_CASE_   :      List[str]         =  False
				SCREAMING_SNAKE_CASE_   :      List[Any]         =  False

				def 							lowercase_     (	self							:				Dict)->   Union[str, Any]:



						'''simple docstring'''


						__lowerCAmelCase:       int          =			OpenLlamaModelTester(self)
						__lowerCAmelCase:       Optional[Any]          =			ConfigTester(self     ,  config_class=UpperCamelCase__     ,  hidden_size=3_7)

				def 							lowercase_     (	self							:				List[str])->   Union[str, Any]:



						'''simple docstring'''


						self.config_tester.run_common_tests()

				def 							lowercase_     (	self							:				Union[str, Any])->   List[Any]:



						'''simple docstring'''


						__lowerCAmelCase:       Any          =			self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*UpperCamelCase__)

				def 							lowercase_     (	self							:				int)->   str:



						'''simple docstring'''


						__lowerCAmelCase:       Optional[Any]          =			self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
								__lowerCAmelCase:       Union[str, Any]          =			type
								self.model_tester.create_and_check_model(*UpperCamelCase__)

				def 							lowercase_     (	self							:				Tuple)->   Optional[Any]:



						'''simple docstring'''


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

				def 							lowercase_     (	self							:				Dict)->   Tuple:



						'''simple docstring'''


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

				def 							lowercase_     (	self							:				Optional[int])->   Any:



						'''simple docstring'''


						__lowerCAmelCase      , __lowerCAmelCase:       str          =			self.model_tester.prepare_config_and_inputs_for_common()
						__lowerCAmelCase:       Tuple          =			3
						__lowerCAmelCase:       Any          =			"multi_label_classification"
						__lowerCAmelCase:       str          =			input_dict["input_ids"]
						__lowerCAmelCase:       Optional[int]          =			input_ids.ne(1).to(UpperCamelCase__)
						__lowerCAmelCase:       Optional[int]          =			ids_tensor(
						    [self.model_tester.batch_size, config.num_labels]     ,  self.model_tester.type_sequence_label_size).to(torch.float)
						__lowerCAmelCase:       Dict          =			OpenLlamaForSequenceClassification(UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase:       Union[str, Any]          =			model(UpperCamelCase__     ,  attention_mask=UpperCamelCase__     ,  labels=UpperCamelCase__)
						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 							lowercase_     (	self							:				Any)->   Tuple:



						'''simple docstring'''


						pass




				@parameterized.expand([("linear",), ("dynamic",)])
				def 							lowercase_     (	self							:				Any     ,  UpperCamelCase__							:				List[str])->   Dict:



						'''simple docstring'''


						__lowerCAmelCase      , __lowerCAmelCase:       Optional[Any]          =			self.model_tester.prepare_config_and_inputs_for_common()
						__lowerCAmelCase:       Any          =			ids_tensor([1, 1_0]     ,  config.vocab_size)
						__lowerCAmelCase:       Union[str, Any]          =			ids_tensor([1, int(config.max_position_embeddings * 1.5)]     ,  config.vocab_size)

						set_seed(4_2)  # Fixed seed at init time so the two models get the same random weights
						__lowerCAmelCase:       List[Any]          =			OpenLlamaModel(UpperCamelCase__)
						original_model.to(UpperCamelCase__)
						original_model.eval()
						__lowerCAmelCase:       int          =			original_model(UpperCamelCase__).last_hidden_state
						__lowerCAmelCase:       str          =			original_model(UpperCamelCase__).last_hidden_state

						set_seed(4_2)  # Fixed seed at init time so the two models get the same random weights
						__lowerCAmelCase:       Dict          =			{"type": scaling_type, "factor": 10.0}
						__lowerCAmelCase:       List[str]          =			OpenLlamaModel(UpperCamelCase__)
						scaled_model.to(UpperCamelCase__)
						scaled_model.eval()
						__lowerCAmelCase:       Dict          =			scaled_model(UpperCamelCase__).last_hidden_state
						__lowerCAmelCase:       Any          =			scaled_model(UpperCamelCase__).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(UpperCamelCase__     ,  UpperCamelCase__     ,  atol=1e-5))
						else:
								self.assertFalse(torch.allclose(UpperCamelCase__     ,  UpperCamelCase__     ,  atol=1e-5))

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


