


import argparse
import gc
import json
import os

import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.utils.deepspeed import DummyOptim, DummyScheduler


a_ :Optional[Any]  							=			16
a_ :Dict  							=			32




def    a   (							A__							)	-> Dict:




      '''simple docstring'''

      return int(x / 2**2_0							)

class 							lowercase  :







     def __enter__(						self :						List[str]      ):
           gc.collect()
           torch.cuda.empty_cache()
           torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
           SCREAMING_SNAKE_CASE__       :  Any						     =		torch.cuda.memory_allocated()
           return self







     def __exit__(						self :						Union[str, Any]   ,   *_lowercase :						Tuple      ):
           gc.collect()
           torch.cuda.empty_cache()
           SCREAMING_SNAKE_CASE__       :  Tuple						     =		torch.cuda.memory_allocated()
           SCREAMING_SNAKE_CASE__       :  Tuple						     =		torch.cuda.max_memory_allocated()
           SCREAMING_SNAKE_CASE__       :  Union[str, Any]						     =		bamb(self.end - self.begin      )
           SCREAMING_SNAKE_CASE__       :  str						     =		bamb(self.peak - self.begin      )
           # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")




def    a   (							A__      ,    A__ = 1_6      ,    A__ = "bert-base-cased"      ,    A__ = 3_2_0      ,    A__ = 1_6_0      ,    )	-> Optional[int]:




      '''simple docstring'''

      SCREAMING_SNAKE_CASE__       :  int						     =		AutoTokenizer.from_pretrained(A__							)
      SCREAMING_SNAKE_CASE__       :  Optional[int]						     =		load_dataset(
          '''glue'''      ,    '''mrpc'''      ,    split={'''train''': f"""train[:{n_train}]""", '''validation''': f"""validation[:{n_val}]"""}							)

      def tokenize_function(A__							):
            # max_length=None => use the model max length (it's actually the default)
            SCREAMING_SNAKE_CASE__       :  List[str]						     =		tokenizer(examples['''sentence1''']      ,    examples['''sentence2''']      ,    truncation=A__      ,    max_length=A__							)
            return outputs

      # Apply the method we just defined to all the examples in all the splits of the dataset
      SCREAMING_SNAKE_CASE__       :  str						     =		datasets.map(
          A__      ,    batched=A__      ,    remove_columns=['''idx''', '''sentence1''', '''sentence2''']      ,    load_from_cache_file=A__							)

      # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
      # transformers library
      SCREAMING_SNAKE_CASE__       :  List[str]						     =		tokenized_datasets.rename_column('''label'''      ,    '''labels'''							)

      def collate_fn(A__							):
            # On TPU it's best to pad everything to the same length or training will be very slow.
            if accelerator.distributed_type == DistributedType.TPU:
                  return tokenizer.pad(A__      ,    padding='''max_length'''      ,    max_length=1_2_8      ,    return_tensors='''pt'''							)
            return tokenizer.pad(A__      ,    padding='''longest'''      ,    return_tensors='''pt'''							)

      # Instantiate dataloaders.
      SCREAMING_SNAKE_CASE__       :  Union[str, Any]						     =		DataLoader(
          tokenized_datasets['''train''']      ,    shuffle=A__      ,    collate_fn=A__      ,    batch_size=A__							)
      SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		DataLoader(
          tokenized_datasets['''validation''']      ,    shuffle=A__      ,    collate_fn=A__      ,    batch_size=A__							)

      return train_dataloader, eval_dataloader




def    a   (							A__      ,    A__							)	-> Union[str, Any]:




      '''simple docstring'''

      SCREAMING_SNAKE_CASE__       :  Union[str, Any]						     =		Accelerator()

      # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
      SCREAMING_SNAKE_CASE__       :  int						     =		config['''lr''']
      SCREAMING_SNAKE_CASE__       :  Optional[int]						     =		int(config['''num_epochs''']							)
      SCREAMING_SNAKE_CASE__       :  List[str]						     =		int(config['''seed''']							)
      SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		int(config['''batch_size''']							)
      SCREAMING_SNAKE_CASE__       :  Tuple						     =		args.model_name_or_path

      set_seed(A__							)
      SCREAMING_SNAKE_CASE__     ,       SCREAMING_SNAKE_CASE__       :  str						     =		get_dataloaders(A__      ,    A__      ,    A__      ,    args.n_train      ,    args.n_val							)

      # Instantiate the model (we build the model here so that the seed also control new weights initialization)
      SCREAMING_SNAKE_CASE__       :  str						     =		AutoModelForSequenceClassification.from_pretrained(A__      ,    return_dict=A__							)

      # Instantiate optimizer
      SCREAMING_SNAKE_CASE__       :  Optional[int]						     =		(
          AdamW
          if accelerator.state.deepspeed_plugin is None
          or '''optimizer''' not in accelerator.state.deepspeed_plugin.deepspeed_config
          else DummyOptim
      )
      SCREAMING_SNAKE_CASE__       :  Dict						     =		optimizer_cls(params=model.parameters()      ,    lr=A__							)

      if accelerator.state.deepspeed_plugin is not None:
            SCREAMING_SNAKE_CASE__       :  Union[str, Any]						     =		accelerator.state.deepspeed_plugin.deepspeed_config[
                '''gradient_accumulation_steps'''
            ]
      else:
            SCREAMING_SNAKE_CASE__       :  Optional[int]						     =		1
      SCREAMING_SNAKE_CASE__       :  Tuple						     =		(len(A__							) * num_epochs) // gradient_accumulation_steps

      # Instantiate scheduler
      if (
          accelerator.state.deepspeed_plugin is None
          or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
      ):
            SCREAMING_SNAKE_CASE__       :  List[str]						     =		get_linear_schedule_with_warmup(
                optimizer=A__      ,    num_warmup_steps=0      ,    num_training_steps=A__      ,    )
      else:
            SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		DummyScheduler(A__      ,    total_num_steps=A__      ,    warmup_num_steps=0							)

      # Prepare everything
      # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
      # prepare method.
      SCREAMING_SNAKE_CASE__     ,       SCREAMING_SNAKE_CASE__     ,       SCREAMING_SNAKE_CASE__     ,       SCREAMING_SNAKE_CASE__     ,       SCREAMING_SNAKE_CASE__       :  Tuple						     =		accelerator.prepare(
          A__      ,    A__      ,    A__      ,    A__      ,    A__							)

      # We need to keep track of how many total steps we have iterated over
      SCREAMING_SNAKE_CASE__       :  Any						     =		0
      # We also need to keep track of the stating epoch so files are named properly
      SCREAMING_SNAKE_CASE__       :  int						     =		0

      # Now we train the model
      SCREAMING_SNAKE_CASE__       :  Dict						     =		{}
      for epoch in range(A__      ,    A__							):
            with TorchTracemalloc() as tracemalloc:
                  model.train()
                  for step, batch in enumerate(A__							):
                        SCREAMING_SNAKE_CASE__       :  List[Any]						     =		model(**A__							)
                        SCREAMING_SNAKE_CASE__       :  Union[str, Any]						     =		outputs.loss
                        SCREAMING_SNAKE_CASE__       :  Tuple						     =		loss / gradient_accumulation_steps
                        accelerator.backward(A__							)
                        if step % gradient_accumulation_steps == 0:
                              optimizer.step()
                              lr_scheduler.step()
                              optimizer.zero_grad()

                        overall_step += 1

        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
            accelerator.print('''Memory before entering the train : {}'''.format(bamb(tracemalloc.begin							)							)							)
            accelerator.print('''Memory consumed at the end of the train (end-begin): {}'''.format(tracemalloc.used							)							)
            accelerator.print('''Peak Memory consumed during the train (max-begin): {}'''.format(tracemalloc.peaked							)							)
            accelerator.print(
                '''Total Peak Memory consumed during the train (max): {}'''.format(
                    tracemalloc.peaked + bamb(tracemalloc.begin							)							)							)
            SCREAMING_SNAKE_CASE__       :  Dict						     =		tracemalloc.peaked + bamb(tracemalloc.begin							)
            if args.peak_memory_upper_bound is not None:
                  assert (
                      train_total_peak_memory[f"""epoch-{epoch}"""] <= args.peak_memory_upper_bound
                  ), "Peak memory usage exceeded the upper bound"

      accelerator.wait_for_everyone()
      if accelerator.is_main_process:
            with open(os.path.join(args.output_dir      ,    '''peak_memory_utilization.json'''							)      ,    '''w'''							) as f:
                  json.dump(A__      ,    A__							)




def    a   (							)	-> Union[str, Any]:




      '''simple docstring'''

      SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		argparse.ArgumentParser(description='''Simple example of training script tracking peak GPU memory usage.'''							)
      parser.add_argument(
          '''--model_name_or_path'''      ,    type=A__      ,    default='''bert-base-cased'''      ,    help='''Path to pretrained model or model identifier from huggingface.co/models.'''      ,    required=A__      ,    )
      parser.add_argument(
          '''--output_dir'''      ,    type=A__      ,    default='''.'''      ,    help='''Optional save directory where all checkpoint folders will be stored. Default is the current working directory.'''      ,    )
      parser.add_argument(
          '''--peak_memory_upper_bound'''      ,    type=A__      ,    default=A__      ,    help='''The upper bound of peak memory usage in MB. If set, the training will throw an error if the peak memory usage exceeds this value.'''      ,    )
      parser.add_argument(
          '''--n_train'''      ,    type=A__      ,    default=3_2_0      ,    help='''Number of training examples to use.'''      ,    )
      parser.add_argument(
          '''--n_val'''      ,    type=A__      ,    default=1_6_0      ,    help='''Number of validation examples to use.'''      ,    )
      parser.add_argument(
          '''--num_epochs'''      ,    type=A__      ,    default=1      ,    help='''Number of train epochs.'''      ,    )
      SCREAMING_SNAKE_CASE__       :  List[str]						     =		parser.parse_args()
      SCREAMING_SNAKE_CASE__       :  Dict						     =		{'''lr''': 2e-5, '''num_epochs''': args.num_epochs, '''seed''': 4_2, '''batch_size''': 1_6}
      training_function(A__      ,    A__							)


if __name__ == "__main__":
    main()









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    )    )


