


'''simple docstring'''
import logging

import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


UpperCAmelCase_ :     Dict								=      logging.getLogger(__name__)

class 						_lowerCamelCase					(    snake_case_				):




 '''simple docstring'''


 def  snake_case__   (      self      ,      __lowercase      ,      __lowercase      ,      __lowercase=None      ,      __lowercase=None						):

  """simple docstring"""







  __A  :		Any       =			self.layer[current_layer](__lowercase      ,      __lowercase      ,      head_mask[current_layer]						)

  __A  :		List[Any]       =			layer_outputs[0]

  return hidden_states


@add_start_docstrings(
    '''The bare Bert Model transformer with PABEE outputting raw hidden-states without any specific head on top.'''    ,		snake_case_    ,		)
class 						_lowerCamelCase					(    snake_case_				):




 '''simple docstring'''


 def __init__(      self      ,      __lowercase						):

  """simple docstring"""







  super().__init__(__lowercase						)

  __A  :		str       =			BertEncoderWithPabee(__lowercase						)

  self.init_weights()
  __A  :		List[Any]       =			0
  __A  :		str       =			0
  __A  :		List[str]       =			0

  __A  :		Tuple       =			0


 def  snake_case__   (      self      ,      __lowercase						):

  """simple docstring"""







  __A  :		Any       =			threshold


 def  snake_case__   (      self      ,      __lowercase						):

  """simple docstring"""







  __A  :		Optional[int]       =			patience


 def  snake_case__   (      self						):

  """simple docstring"""







  __A  :		Optional[Any]       =			0
  __A  :		Union[str, Any]       =			0


 def  snake_case__   (      self						):

  """simple docstring"""







  __A  :		Optional[int]       =			self.inference_layers_num / self.inference_instances_num
  __A  :		Optional[int]       =			(
      F"""*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ="""
      F""" {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***"""
  )
  print(__lowercase						)







 @add_start_docstrings_to_model_forward(__lowercase						)
 def  snake_case__   (      self      ,      __lowercase=None      ,      __lowercase=None      ,      __lowercase=None      ,      __lowercase=None      ,      __lowercase=None      ,      __lowercase=None      ,      __lowercase=None      ,      __lowercase=None      ,      __lowercase=None      ,      __lowercase=None      ,      __lowercase=False      ,      ):

  """simple docstring"""







  if input_ids is not None and inputs_embeds is not None:
   raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time'						)
  elif input_ids is not None:
   __A  :		int       =			input_ids.size()
  elif inputs_embeds is not None:
   __A  :		Optional[int]       =			inputs_embeds.size()[:-1]
  else:
   raise ValueError('You have to specify either input_ids or inputs_embeds'						)

  __A  :		List[Any]       =			input_ids.device if input_ids is not None else inputs_embeds.device

  if attention_mask is None:
   __A  :		Optional[int]       =			torch.ones(__lowercase      ,      device=__lowercase						)
  if token_type_ids is None:
   __A  :		Optional[Any]       =			torch.zeros(__lowercase      ,      dtype=torch.long      ,      device=__lowercase						)

  # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
  # ourselves in which case we just need to make it broadcastable to all heads.
  __A  :		torch.Tensor       =			self.get_extended_attention_mask(__lowercase      ,      __lowercase      ,      __lowercase						)

  # If a 2D ou 3D attention mask is provided for the cross-attention
  # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
  if self.config.is_decoder and encoder_hidden_states is not None:
   __A     ,__A     ,__A  :		Dict       =			encoder_hidden_states.size()
   __A  :		Dict       =			(encoder_batch_size, encoder_sequence_length)
   if encoder_attention_mask is None:
    __A  :		Optional[Any]       =			torch.ones(__lowercase      ,      device=__lowercase						)
   __A  :		List[Any]       =			self.invert_attention_mask(__lowercase						)
  else:
   __A  :		Any       =			None

  # Prepare head mask if needed
  # 1.0 in head_mask indicate we keep the head
  # attention_probs has shape bsz x n_heads x N x N
  # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
  # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
  __A  :		List[str]       =			self.get_head_mask(__lowercase      ,      self.config.num_hidden_layers						)

  __A  :		Tuple       =			self.embeddings(
      input_ids=__lowercase      ,      position_ids=__lowercase      ,      token_type_ids=__lowercase      ,      inputs_embeds=__lowercase						)
  __A  :		Union[str, Any]       =			embedding_output

  if self.training:
   __A  :		Optional[Any]       =			[]
   for i in range(self.config.num_hidden_layers						):
    __A  :		Optional[Any]       =			self.encoder.adaptive_forward(
        __lowercase      ,      current_layer=__lowercase      ,      attention_mask=__lowercase      ,      head_mask=__lowercase						)

    __A  :		Union[str, Any]       =			self.pooler(__lowercase						)
    __A  :		Dict       =			output_layers[i](output_dropout(__lowercase						)						)
    res.append(__lowercase						)
  elif self.patience == 0:  # Use all layers for inference
   __A  :		Dict       =			self.encoder(
       __lowercase      ,      attention_mask=__lowercase      ,      head_mask=__lowercase      ,      encoder_hidden_states=__lowercase      ,      encoder_attention_mask=__lowercase      ,      )
   __A  :		List[str]       =			self.pooler(encoder_outputs[0]						)
   __A  :		Dict       =			[output_layers[self.config.num_hidden_layers - 1](__lowercase						)]
  else:
   __A  :		Dict       =			0
   __A  :		int       =			None
   __A  :		Any       =			0
   for i in range(self.config.num_hidden_layers						):
    calculated_layer_num += 1
    __A  :		Dict       =			self.encoder.adaptive_forward(
        __lowercase      ,      current_layer=__lowercase      ,      attention_mask=__lowercase      ,      head_mask=__lowercase						)

    __A  :		Optional[Any]       =			self.pooler(__lowercase						)
    __A  :		str       =			output_layers[i](__lowercase						)
    if regression:
     __A  :		Optional[Any]       =			logits.detach()
     if patient_result is not None:
      __A  :		Dict       =			patient_result.detach()
     if (patient_result is not None) and torch.abs(patient_result - labels						) < self.regression_threshold:
      patient_counter += 1
     else:
      __A  :		Union[str, Any]       =			0
    else:
     __A  :		Tuple       =			logits.detach().argmax(dim=1						)
     if patient_result is not None:
      __A  :		Any       =			patient_result.detach().argmax(dim=1						)
     if (patient_result is not None) and torch.all(labels.eq(__lowercase						)						):
      patient_counter += 1
     else:
      __A  :		Optional[int]       =			0

    __A  :		Tuple       =			logits
    if patient_counter == self.patience:
     break
   __A  :		Optional[Any]       =			[patient_result]
   self.inference_layers_num += calculated_layer_num
   self.inference_instances_num += 1

  return res


@add_start_docstrings(
    '''Bert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of
    the pooled output) e.g. for GLUE tasks. '''    ,		snake_case_    ,		)
class 						_lowerCamelCase					(    snake_case_				):




 '''simple docstring'''


 def __init__(      self      ,      __lowercase						):

  """simple docstring"""







  super().__init__(__lowercase						)
  __A  :		List[str]       =			config.num_labels

  __A  :		Union[str, Any]       =			BertModelWithPabee(__lowercase						)
  __A  :		Dict       =			nn.Dropout(config.hidden_dropout_prob						)
  __A  :		Optional[Any]       =			nn.ModuleList(
      [nn.Linear(config.hidden_size      ,      self.config.num_labels						) for _ in range(config.num_hidden_layers						)]						)

  self.init_weights()







 @add_start_docstrings_to_model_forward(__lowercase						)
 def  snake_case__   (      self      ,      __lowercase=None      ,      __lowercase=None      ,      __lowercase=None      ,      __lowercase=None      ,      __lowercase=None      ,      __lowercase=None      ,      __lowercase=None      ,      ):

  """simple docstring"""







  __A  :		List[str]       =			self.bert(
      input_ids=__lowercase      ,      attention_mask=__lowercase      ,      token_type_ids=__lowercase      ,      position_ids=__lowercase      ,      head_mask=__lowercase      ,      inputs_embeds=__lowercase      ,      output_dropout=self.dropout      ,      output_layers=self.classifiers      ,      regression=self.num_labels == 1      ,      )

  __A  :		Tuple       =			(logits[-1],)

  if labels is not None:
   __A  :		str       =			None
   __A  :		Dict       =			0
   for ix, logits_item in enumerate(__lowercase						):
    if self.num_labels == 1:
     #  We are doing regression
     __A  :		Optional[int]       =			MSELoss()
     __A  :		Tuple       =			loss_fct(logits_item.view(-1						)      ,      labels.view(-1						)						)
    else:
     __A  :		List[str]       =			CrossEntropyLoss()
     __A  :		List[str]       =			loss_fct(logits_item.view(-1      ,      self.num_labels						)      ,      labels.view(-1						)						)
    if total_loss is None:
     __A  :		str       =			loss
    else:
     total_loss += loss * (ix + 1)
    total_weights += ix + 1
   __A  :		List[Any]       =			(total_loss / total_weights,) + outputs

  return outputs










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




 '''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      ,      ):

  """simple docstring"""







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


 def  snake_case__   (      self						):

  """simple docstring"""







  __A  :		Union[str, Any]       =			ids_tensor([self.batch_size, self.seq_length]      ,      self.vocab_size						)

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

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

  __A  :		Union[str, Any]       =			None
  __A  :		Optional[int]       =			None
  __A  :		List[str]       =			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  :		Union[str, Any]       =			ids_tensor([self.batch_size]      ,      self.num_choices						)

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

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


 def  snake_case__   (      self						):

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


 def  snake_case__   (      self      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase						):

  """simple docstring"""







  __A  :		List[str]       =			OpenLlamaModel(config=__lowercase						)
  model.to(__lowercase						)
  model.eval()
  __A  :		Any       =			model(__lowercase      ,      attention_mask=__lowercase						)
  __A  :		Dict       =			model(__lowercase						)
  self.parent.assertEqual(result.last_hidden_state.shape      ,      (self.batch_size, self.seq_length, self.hidden_size)						)


 def  snake_case__   (      self      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      ):

  """simple docstring"""







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


 def  snake_case__   (      self      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      ):

  """simple docstring"""







  __A  :		Dict       =			OpenLlamaForCausalLM(config=__lowercase						)
  model.to(__lowercase						)
  model.eval()
  __A  :		Optional[Any]       =			model(__lowercase      ,      attention_mask=__lowercase      ,      labels=__lowercase						)
  self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.vocab_size)						)


 def  snake_case__   (      self      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      __lowercase      ,      ):

  """simple docstring"""







  __A  :		List[Any]       =			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  :		Dict       =			ids_tensor((self.batch_size, 3)      ,      config.vocab_size						)
  __A  :		Optional[Any]       =			ids_tensor((self.batch_size, 3)      ,      vocab_size=2						)

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

  __A  :		Tuple       =			model(
      __lowercase      ,      attention_mask=__lowercase      ,      encoder_hidden_states=__lowercase      ,      encoder_attention_mask=__lowercase      ,      output_hidden_states=__lowercase      ,      )['hidden_states'][0]
  __A  :		Any       =			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  :		Any       =			ids_tensor((1,)      ,      output_from_past.shape[-1]						).item()
  __A  :		Union[str, Any]       =			output_from_no_past[:, -3:, random_slice_idx].detach()
  __A  :		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  snake_case__   (      self						):

  """simple docstring"""







  __A  :		int       =			self.prepare_config_and_inputs()
  (
      (
      __A
  )     ,(
      __A
  )     ,(
      __A
  )     ,(
      __A
  )     ,(
      __A
  )     ,(
      __A
  )     ,(
      __A
  )     ,
  )  :		Optional[Any]       =			config_and_inputs
  __A  :		Optional[int]       =			{'input_ids': input_ids, 'attention_mask': input_mask}
  return config, inputs_dict


@require_torch
class 						_lowerCamelCase					(    snake_case_    ,		snake_case_    ,		snake_case_    ,		unittest.TestCase				):




 '''simple docstring'''
 __lowercase				:		List[str]    							=						(
     (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
 )
 __lowercase				:		List[Any]    							=						(OpenLlamaForCausalLM,) if is_torch_available() else ()
 __lowercase				:		Optional[int]    							=						(
     {
         '''feature-extraction''': OpenLlamaModel,
         '''text-classification''': OpenLlamaForSequenceClassification,
         '''text-generation''': OpenLlamaForCausalLM,
         '''zero-shot''': OpenLlamaForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 __lowercase				:		Union[str, Any]    							=						False
 __lowercase				:		Union[str, Any]    							=						False


 def  snake_case__   (      self						):

  """simple docstring"""







  __A  :		Optional[int]       =			OpenLlamaModelTester(self						)
  __A  :		List[str]       =			ConfigTester(self      ,      config_class=__lowercase      ,      hidden_size=37						)


 def  snake_case__   (      self						):

  """simple docstring"""







  self.config_tester.run_common_tests()


 def  snake_case__   (      self						):

  """simple docstring"""







  __A  :		int       =			self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_model(*__lowercase						)


 def  snake_case__   (      self						):

  """simple docstring"""







  __A  :		int       =			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(*__lowercase						)


 def  snake_case__   (      self						):

  """simple docstring"""







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


 def  snake_case__   (      self						):

  """simple docstring"""







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


 def  snake_case__   (      self						):

  """simple docstring"""







  __A     ,__A  :		Optional[int]       =			self.model_tester.prepare_config_and_inputs_for_common()
  __A  :		str       =			3
  __A  :		int       =			'multi_label_classification'
  __A  :		Union[str, Any]       =			input_dict['input_ids']
  __A  :		str       =			input_ids.ne(1						).to(__lowercase						)
  __A  :		Union[str, Any]       =			ids_tensor(
      [self.model_tester.batch_size, config.num_labels]      ,      self.model_tester.type_sequence_label_size						).to(torch.float						)
  __A  :		int       =			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)						)


 @unittest.skip('Open-Llama buffers include complex numbers, which breaks this test'						)
 def  snake_case__   (      self						):

  """simple docstring"""







  pass







 @parameterized.expand([('linear',), ('dynamic',)]						)
 def  snake_case__   (      self      ,      __lowercase						):

  """simple docstring"""







  __A     ,__A  :		str       =			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  :		Optional[int]       =			original_model(__lowercase						).last_hidden_state
  __A  :		int       =			original_model(__lowercase						).last_hidden_state

  set_seed(42						)  # Fixed seed at init time so the two models get the same random weights
  __A  :		List[Any]       =			{'type': scaling_type, 'factor': 1_0.0}
  __A  :		str       =			OpenLlamaModel(__lowercase						)
  scaled_model.to(__lowercase						)
  scaled_model.eval()
  __A  :		Dict       =			scaled_model(__lowercase						).last_hidden_state
  __A  :		List[str]       =			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						)						)






