

"""simple docstring"""


from copy import deepcopy

import torch
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader

from accelerate.accelerator import Accelerator
from accelerate.state import GradientState
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import DistributedType, is_torch_version, set_seed






def     a__    (    __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )	->	Any:
		for param, grad_param in zip(model_a.parameters()					, model_b.parameters() ):
				if not param.requires_grad:
						continue
				if not did_step:
						# Grads should not be in sync
						assert (
						    torch.allclose(param.grad					, grad_param.grad ) is False
						), F"Gradients in sync when they should not be at iteration {iteration}:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})"
				else:
						# Grads should be in sync
						assert (
						    torch.allclose(param.grad					, grad_param.grad ) is True
						), F"Gradients not in sync when they should be at iteration {iteration}:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})"






def     a__    (    __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE=True )	->	str:
		model.train()
		__lowerCAmelCase:       Any          =			model(__SCREAMING_SNAKE_CASE )
		__lowerCAmelCase:       Optional[int]          =			F.mse_loss(__SCREAMING_SNAKE_CASE					, target.to(output.device ) )
		if not do_backward:
				loss /= accelerator.gradient_accumulation_steps
				loss.backward()
		else:
				accelerator.backward(__SCREAMING_SNAKE_CASE )






def     a__    (    __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE=False )	->	Tuple:
		set_seed(4_2 )
		__lowerCAmelCase:       Any          =			RegressionModel()
		__lowerCAmelCase:       int          =			deepcopy(__SCREAMING_SNAKE_CASE )
		__lowerCAmelCase:       str          =			RegressionDataset(length=8_0 )
		__lowerCAmelCase:       Optional[Any]          =			DataLoader(__SCREAMING_SNAKE_CASE					, batch_size=1_6 )
		model.to(accelerator.device )
		if sched:
				__lowerCAmelCase:       Optional[Any]          =			AdamW(params=model.parameters()					, lr=1E-3 )
				__lowerCAmelCase:       Tuple          =			AdamW(params=ddp_model.parameters()					, lr=1E-3 )
				__lowerCAmelCase:       Tuple          =			LambdaLR(__SCREAMING_SNAKE_CASE					, lr_lambda=lambda __SCREAMING_SNAKE_CASE : epoch**0.65 )
				__lowerCAmelCase:       Dict          =			LambdaLR(__SCREAMING_SNAKE_CASE					, lr_lambda=lambda __SCREAMING_SNAKE_CASE : epoch**0.65 )
		# Make a copy of `model`
		if sched:
				__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase:       Any          =			accelerator.prepare(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )
		else:
				__lowerCAmelCase      , __lowerCAmelCase:       Dict          =			accelerator.prepare(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )
		if sched:
				return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched)
		return model, ddp_model, dataloader






def     a__    (    __SCREAMING_SNAKE_CASE )	->	List[str]:
		# Test when on a single CPU or GPU that the context manager does nothing
		__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase:       Union[str, Any]          =			get_training_setup(__SCREAMING_SNAKE_CASE )
		# Use a single batch
		__lowerCAmelCase      , __lowerCAmelCase:       Optional[Any]          =			next(iter(__SCREAMING_SNAKE_CASE ) ).values()
		for iteration in range(3 ):
				# Gather the distributed inputs and targs for the base model
				__lowerCAmelCase      , __lowerCAmelCase:       List[str]          =			accelerator.gather((ddp_input, ddp_target) )
				__lowerCAmelCase      , __lowerCAmelCase:       Dict          =			input.to(accelerator.device ), target.to(accelerator.device )
				# Perform our initial ground truth step in non "DDP"
				step_model(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )
				# Do "gradient accumulation" (noop)
				if iteration % 2 == 0:
						# Accumulate grads locally
						with accelerator.no_sync(__SCREAMING_SNAKE_CASE ):
								step_model(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )
				else:
						# Sync grads
						step_model(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )

				# Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
				check_model_parameters(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )
				for param, ddp_param in zip(model.parameters()					, ddp_model.parameters() ):
						if not param.requires_grad:
								continue
						assert torch.allclose(
						    param.grad					, ddp_param.grad ), F"Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"

				# Shuffle ddp_input on each iteration
				torch.manual_seed(1_3_3_7 + iteration )
				__lowerCAmelCase:       Dict          =			ddp_input[torch.randperm(len(__SCREAMING_SNAKE_CASE ) )]






def     a__    (    __SCREAMING_SNAKE_CASE )	->	Any:
		# Test on distributed setup that context manager behaves properly
		__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase:       Tuple          =			get_training_setup(__SCREAMING_SNAKE_CASE )
		# Use a single batch
		__lowerCAmelCase      , __lowerCAmelCase:       str          =			next(iter(__SCREAMING_SNAKE_CASE ) ).values()
		for iteration in range(3 ):
				# Gather the distributed inputs and targs for the base model
				__lowerCAmelCase      , __lowerCAmelCase:       Dict          =			accelerator.gather((ddp_input, ddp_target) )
				__lowerCAmelCase      , __lowerCAmelCase:       Tuple          =			input.to(accelerator.device ), target.to(accelerator.device )
				# Perform our initial ground truth step in non "DDP"
				step_model(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )
				# Do "gradient accumulation" (noop)
				if iteration % 2 == 0:
						# Accumulate grads locally
						with accelerator.no_sync(__SCREAMING_SNAKE_CASE ):
								step_model(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )
				else:
						# Sync grads
						step_model(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )

				# DDP model and model should only be in sync when not (iteration % 2 == 0)
				for param, ddp_param in zip(model.parameters()					, ddp_model.parameters() ):
						if not param.requires_grad:
								continue
						if iteration % 2 == 0:
								# Grads should not be in sync
								assert (
								    torch.allclose(param.grad					, ddp_param.grad ) is False
								), F"Gradients in sync when they should not be:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})"
						else:
								# Grads should be in sync
								assert (
								    torch.allclose(param.grad					, ddp_param.grad ) is True
								), F"Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"

        # Shuffle ddp_input on each iteration
				torch.manual_seed(1_3_3_7 + iteration )
				__lowerCAmelCase:       Dict          =			ddp_input[torch.randperm(len(__SCREAMING_SNAKE_CASE ) )]






def     a__    (    __SCREAMING_SNAKE_CASE=False					, __SCREAMING_SNAKE_CASE=False )	->	int:
		__lowerCAmelCase:       str          =			Accelerator(
		    split_batches=__SCREAMING_SNAKE_CASE					, dispatch_batches=__SCREAMING_SNAKE_CASE					, gradient_accumulation_steps=2 )
		# Test that context manager behaves properly
		__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase:       Dict          =			get_training_setup(__SCREAMING_SNAKE_CASE )
		for iteration, batch in enumerate(__SCREAMING_SNAKE_CASE ):
				__lowerCAmelCase      , __lowerCAmelCase:       Dict          =			batch.values()
				# Gather the distributed inputs and targs for the base model
				__lowerCAmelCase      , __lowerCAmelCase:       List[Any]          =			accelerator.gather((ddp_input, ddp_target) )
				__lowerCAmelCase      , __lowerCAmelCase:       Any          =			input.to(accelerator.device ), target.to(accelerator.device )
				# Perform our initial ground truth step in non "DDP"
				step_model(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )
				# Do "gradient accumulation" (noop)
				with accelerator.accumulate(__SCREAMING_SNAKE_CASE ):
						step_model(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )

				# DDP model and model should only be in sync when not (iteration % 2 == 0)
				for param, ddp_param in zip(model.parameters()					, ddp_model.parameters() ):
						if not param.requires_grad:
								continue
						if ((iteration + 1) % 2 == 0) or (iteration == len(__SCREAMING_SNAKE_CASE ) - 1):
								# Grads should be in sync
								assert (
								    torch.allclose(param.grad					, ddp_param.grad ) is True
								), F"Gradients not in sync when they should be at iteration {iteration}:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"
						else:
								# Grads should not be in sync
								assert (
								    torch.allclose(param.grad					, ddp_param.grad ) is False
								), F"Gradients in sync when they should not be at iteration {iteration}:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})"

        # Shuffle ddp_input on each iteration
				torch.manual_seed(1_3_3_7 + iteration )
				__lowerCAmelCase:       Tuple          =			ddp_input[torch.randperm(len(__SCREAMING_SNAKE_CASE ) )]
		GradientState._reset_state()






def     a__    (    __SCREAMING_SNAKE_CASE=False					, __SCREAMING_SNAKE_CASE=False )	->	Optional[Any]:
		__lowerCAmelCase:       int          =			Accelerator(
		    split_batches=__SCREAMING_SNAKE_CASE					, dispatch_batches=__SCREAMING_SNAKE_CASE					, gradient_accumulation_steps=2 )
		# Test that context manager behaves properly
		__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase:       Union[str, Any]          =			get_training_setup(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )
		for iteration, batch in enumerate(__SCREAMING_SNAKE_CASE ):
				__lowerCAmelCase      , __lowerCAmelCase:       Optional[int]          =			batch.values()
				# Gather the distributed inputs and targs for the base model
				__lowerCAmelCase      , __lowerCAmelCase:       str          =			accelerator.gather((ddp_input, ddp_target) )
				__lowerCAmelCase      , __lowerCAmelCase:       int          =			input.to(accelerator.device ), target.to(accelerator.device )
				# Perform our initial ground truth step in non "DDP"
				model.train()
				ddp_model.train()
				step_model(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )
				opt.step()

				if ((iteration + 1) % 2 == 0) or ((iteration + 1) == len(__SCREAMING_SNAKE_CASE )):
						if split_batches:
								sched.step()
						else:
								for _ in range(accelerator.num_processes ):
										sched.step()
				opt.zero_grad()
				# Perform gradient accumulation under wrapper
				with accelerator.accumulate(__SCREAMING_SNAKE_CASE ):
						step_model(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )
						ddp_opt.step()
						ddp_sched.step()
						ddp_opt.zero_grad()

				# Learning rates should be the same
				assert (
				    opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"]
				), F"Learning rates found in each optimizer did not align\nopt: {opt.param_groups[0]['lr']}\nDDP opt: {ddp_opt.param_groups[0]['lr']}\n"
				__lowerCAmelCase:       Optional[Any]          =			(((iteration + 1) % 2) == 0) or ((iteration + 1) == len(__SCREAMING_SNAKE_CASE ))
				if accelerator.num_processes > 1:
						check_model_parameters(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )
				# Shuffle ddp_input on each iteration
				torch.manual_seed(1_3_3_7 + iteration )
		GradientState._reset_state()






def     a__    (    )	->	List[str]:
		__lowerCAmelCase:       int          =			Accelerator()

		__lowerCAmelCase:       List[Any]          =			RegressionDataset(length=8_0 )
		__lowerCAmelCase:       Optional[int]          =			DataLoader(__SCREAMING_SNAKE_CASE					, batch_size=1_6 )
		__lowerCAmelCase:       str          =			RegressionDataset(length=9_6 )
		__lowerCAmelCase:       Optional[int]          =			DataLoader(__SCREAMING_SNAKE_CASE					, batch_size=1_6 )
		__lowerCAmelCase      , __lowerCAmelCase:       Union[str, Any]          =			accelerator.prepare(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )
		assert accelerator.gradient_state.active_dataloader is None
		for iteration, _ in enumerate(__SCREAMING_SNAKE_CASE ):
				assert id(accelerator.gradient_state.active_dataloader ) == id(__SCREAMING_SNAKE_CASE )
				if iteration < len(__SCREAMING_SNAKE_CASE ) - 1:
						assert not accelerator.gradient_state.end_of_dataloader
						if iteration == 1:
								for batch_num, _ in enumerate(__SCREAMING_SNAKE_CASE ):
										assert id(accelerator.gradient_state.active_dataloader ) == id(__SCREAMING_SNAKE_CASE )
										if batch_num < len(__SCREAMING_SNAKE_CASE ) - 1:
												assert not accelerator.gradient_state.end_of_dataloader
										else:
												assert accelerator.gradient_state.end_of_dataloader
				else:
						assert accelerator.gradient_state.end_of_dataloader
		assert accelerator.gradient_state.active_dataloader is None






def     a__    (    )	->	Optional[int]:
		__lowerCAmelCase:       Optional[Any]          =			Accelerator()
		__lowerCAmelCase:       Union[str, Any]          =			accelerator.state
		if state.local_process_index == 0:
				print("**Test `accumulate` gradient accumulation with dataloader break**" )
		test_dataloader_break()
		if state.distributed_type == DistributedType.NO:
				if state.local_process_index == 0:
						print("**Test NOOP `no_sync` context manager**" )
				test_noop_sync(__SCREAMING_SNAKE_CASE )
		if state.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_CPU):
				if state.local_process_index == 0:
						print("**Test Distributed `no_sync` context manager**" )
				test_distributed_sync(__SCREAMING_SNAKE_CASE )
		if state.distributed_type == DistributedType.MULTI_GPU:
				for split_batch in [True, False]:
						for dispatch_batches in [True, False]:
								if state.local_process_index == 0:
										print(
										    "**Test `accumulate` gradient accumulation, "					, F"`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**"					, )
								test_gradient_accumulation(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )

    # Currently will break on torch 2.0 +, need to investigate why
		if is_torch_version("<"					, "2.0" ) or state.distributed_type == DistributedType.NO:
				if state.local_process_index == 0:
						print(
						    "**Test `accumulate` gradient accumulation with optimizer and scheduler, "					, "`split_batches=False`, `dispatch_batches=False`**"					, )
				test_gradient_accumulation_with_opt_and_scheduler()
				if state.distributed_type == DistributedType.MULTI_GPU:
						for split_batch in [True, False]:
								for dispatch_batches in [True, False]:
										if not split_batch and not dispatch_batches:
												continue
										if state.local_process_index == 0:
												print(
												    "**Test `accumulate` gradient accumulation with optimizer and scheduler, "					, F"`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**"					, )
										test_gradient_accumulation_with_opt_and_scheduler(__SCREAMING_SNAKE_CASE					, __SCREAMING_SNAKE_CASE )






def     a__    (    __SCREAMING_SNAKE_CASE )	->	Any:
		# For xla_spawn (TPUs)
		main()


if __name__ == "__main__":
			main()





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


