






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				,		_UpperCamelCase				,		_UpperCamelCase=1_3				,		_UpperCamelCase=7				,		_UpperCamelCase=True				,		_UpperCamelCase=True				,		_UpperCamelCase=False				,		_UpperCamelCase=True				,		_UpperCamelCase=9_9				,		_UpperCamelCase=3_2				,		_UpperCamelCase=5				,		_UpperCamelCase=4				,		_UpperCamelCase=3_7				,		_UpperCamelCase="gelu"				,		_UpperCamelCase=0.1				,		_UpperCamelCase=0.1				,		_UpperCamelCase=5_1_2				,		_UpperCamelCase=1_6				,		_UpperCamelCase=2				,		_UpperCamelCase=0.02				,		_UpperCamelCase=3				,		_UpperCamelCase=4				,		_UpperCamelCase=None				,		)	->							Tuple:
							UpperCAmelCase_    :      Any            =						parent
							UpperCAmelCase_    :      List[str]            =						batch_size
							UpperCAmelCase_    :      Optional[Any]            =						seq_length
							UpperCAmelCase_    :      str            =						is_training
							UpperCAmelCase_    :      Optional[int]            =						use_input_mask
							UpperCAmelCase_    :      int            =						use_token_type_ids
							UpperCAmelCase_    :      List[Any]            =						use_labels
							UpperCAmelCase_    :      Dict            =						vocab_size
							UpperCAmelCase_    :      int            =						hidden_size
							UpperCAmelCase_    :      Dict            =						num_hidden_layers
							UpperCAmelCase_    :      Any            =						num_attention_heads
							UpperCAmelCase_    :      List[str]            =						intermediate_size
							UpperCAmelCase_    :      str            =						hidden_act
							UpperCAmelCase_    :      Any            =						hidden_dropout_prob
							UpperCAmelCase_    :      Tuple            =						attention_probs_dropout_prob
							UpperCAmelCase_    :      Union[str, Any]            =						max_position_embeddings
							UpperCAmelCase_    :      Dict            =						type_vocab_size
							UpperCAmelCase_    :      List[str]            =						type_sequence_label_size
							UpperCAmelCase_    :      int            =						initializer_range
							UpperCAmelCase_    :      Dict            =						num_labels
							UpperCAmelCase_    :      Tuple            =						num_choices
							UpperCAmelCase_    :      int            =						scope



		def 					__UpperCAmelCase			(							self   )	->							Union[str, Any]:
							UpperCAmelCase_    :      int            =						ids_tensor([self.batch_size, self.seq_length]				,		self.vocab_size   )

							UpperCAmelCase_    :      str            =						None
							if self.use_input_mask:
												UpperCAmelCase_    :      Dict            =						random_attention_mask([self.batch_size, self.seq_length]   )

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

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

							UpperCAmelCase_    :      Union[str, Any]            =						self.get_config()

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



		def 					__UpperCAmelCase			(							self   )	->							Optional[int]:
							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 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							Tuple:
							UpperCAmelCase_    :      Dict            =						OpenLlamaModel(config=_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.eval()
							UpperCAmelCase_    :      Dict            =						model(_UpperCamelCase				,		attention_mask=_UpperCamelCase   )
							UpperCAmelCase_    :      Dict            =						model(_UpperCamelCase   )
							self.parent.assertEqual(result.last_hidden_state.shape				,		(self.batch_size, self.seq_length, self.hidden_size)   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		)	->							List[Any]:
							UpperCAmelCase_    :      Optional[int]            =						True
							UpperCAmelCase_    :      int            =						OpenLlamaModel(_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.eval()
							UpperCAmelCase_    :      Dict            =						model(
							    _UpperCamelCase				,		attention_mask=_UpperCamelCase				,		encoder_hidden_states=_UpperCamelCase				,		encoder_attention_mask=_UpperCamelCase				,		)
							UpperCAmelCase_    :      List[Any]            =						model(
							    _UpperCamelCase				,		attention_mask=_UpperCamelCase				,		encoder_hidden_states=_UpperCamelCase				,		)
							UpperCAmelCase_    :      Any            =						model(_UpperCamelCase				,		attention_mask=_UpperCamelCase   )
							self.parent.assertEqual(result.last_hidden_state.shape				,		(self.batch_size, self.seq_length, self.hidden_size)   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		)	->							Optional[Any]:
							UpperCAmelCase_    :      Optional[int]            =						OpenLlamaForCausalLM(config=_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.eval()
							UpperCAmelCase_    :      List[str]            =						model(_UpperCamelCase				,		attention_mask=_UpperCamelCase				,		labels=_UpperCamelCase   )
							self.parent.assertEqual(result.logits.shape				,		(self.batch_size, self.seq_length, self.vocab_size)   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		)	->							str:
							UpperCAmelCase_    :      Tuple            =						True
							UpperCAmelCase_    :      str            =						True
							UpperCAmelCase_    :      List[Any]            =						OpenLlamaForCausalLM(config=_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.eval()

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

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

							# append to next input_ids and
							UpperCAmelCase_    :      Optional[int]            =						torch.cat([input_ids, next_tokens]				,		dim=-1   )
							UpperCAmelCase_    :      Union[str, Any]            =						torch.cat([input_mask, next_mask]				,		dim=-1   )

							UpperCAmelCase_    :      Optional[int]            =						model(
							    _UpperCamelCase				,		attention_mask=_UpperCamelCase				,		encoder_hidden_states=_UpperCamelCase				,		encoder_attention_mask=_UpperCamelCase				,		output_hidden_states=_UpperCamelCase				,		)['hidden_states'][0]
							UpperCAmelCase_    :      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
							UpperCAmelCase_    :      Dict            =						ids_tensor((1,)				,		output_from_past.shape[-1]   ).item()
							UpperCAmelCase_    :      Any            =						output_from_no_past[:, -3:, random_slice_idx].detach()
							UpperCAmelCase_    :      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 					__UpperCAmelCase			(							self   )	->							Tuple:
							UpperCAmelCase_    :      Optional[int]            =						self.prepare_config_and_inputs()
							(
							    (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  
							)    :      Any            =						config_and_inputs
							UpperCAmelCase_    :      Union[str, Any]            =						{'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'''



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



		def 					__UpperCAmelCase			(							self   )	->							Optional[Any]:
							UpperCAmelCase_    :      Dict            =						OpenLlamaModelTester(self   )
							UpperCAmelCase_    :      str            =						ConfigTester(self				,		config_class=_UpperCamelCase				,		hidden_size=3_7   )



		def 					__UpperCAmelCase			(							self   )	->							List[str]:
							self.config_tester.run_common_tests()



		def 					__UpperCAmelCase			(							self   )	->							List[Any]:
							UpperCAmelCase_    :      str            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							int:
							UpperCAmelCase_    :      int            =						self.model_tester.prepare_config_and_inputs()
							for type in ["absolute", "relative_key", "relative_key_query"]:
												UpperCAmelCase_    :      Optional[int]            =						type
												self.model_tester.create_and_check_model(*_UpperCamelCase   )



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



		def 					__UpperCAmelCase			(							self   )	->							List[Any]:
							UpperCAmelCase_       ,  UpperCAmelCase_    :      int            =						self.model_tester.prepare_config_and_inputs_for_common()
							UpperCAmelCase_    :      Optional[int]            =						3
							UpperCAmelCase_    :      List[Any]            =						'single_label_classification'
							UpperCAmelCase_    :      int            =						input_dict['input_ids']
							UpperCAmelCase_    :      List[str]            =						input_ids.ne(1   ).to(_UpperCamelCase   )
							UpperCAmelCase_    :      Tuple            =						ids_tensor([self.model_tester.batch_size]				,		self.model_tester.type_sequence_label_size   )
							UpperCAmelCase_    :      Optional[Any]            =						OpenLlamaForSequenceClassification(_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.eval()
							UpperCAmelCase_    :      str            =						model(_UpperCamelCase				,		attention_mask=_UpperCamelCase				,		labels=_UpperCamelCase   )
							self.assertEqual(result.logits.shape				,		(self.model_tester.batch_size, self.model_tester.num_labels)   )



		def 					__UpperCAmelCase			(							self   )	->							Any:
							UpperCAmelCase_       ,  UpperCAmelCase_    :      str            =						self.model_tester.prepare_config_and_inputs_for_common()
							UpperCAmelCase_    :      List[str]            =						3
							UpperCAmelCase_    :      List[str]            =						'multi_label_classification'
							UpperCAmelCase_    :      Tuple            =						input_dict['input_ids']
							UpperCAmelCase_    :      Tuple            =						input_ids.ne(1   ).to(_UpperCamelCase   )
							UpperCAmelCase_    :      str            =						ids_tensor(
							    [self.model_tester.batch_size, config.num_labels]				,		self.model_tester.type_sequence_label_size   ).to(torch.float   )
							UpperCAmelCase_    :      Optional[int]            =						OpenLlamaForSequenceClassification(_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.eval()
							UpperCAmelCase_    :      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 					__UpperCAmelCase			(							self   )	->							Dict:
							pass



		@parameterized.expand([('linear',), ('dynamic',)]   )
		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase   )	->							Any:
							UpperCAmelCase_       ,  UpperCAmelCase_    :      List[str]            =						self.model_tester.prepare_config_and_inputs_for_common()
							UpperCAmelCase_    :      Any            =						ids_tensor([1, 1_0]				,		config.vocab_size   )
							UpperCAmelCase_    :      Dict            =						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
							UpperCAmelCase_    :      str            =						OpenLlamaModel(_UpperCamelCase   )
							original_model.to(_UpperCamelCase   )
							original_model.eval()
							UpperCAmelCase_    :      Any            =						original_model(_UpperCamelCase   ).last_hidden_state
							UpperCAmelCase_    :      Optional[Any]            =						original_model(_UpperCamelCase   ).last_hidden_state

							set_seed(4_2   )  # Fixed seed at init time so the two models get the same random weights
							UpperCAmelCase_    :      Any            =						{'type': scaling_type, 'factor': 10.0}
							UpperCAmelCase_    :      Optional[int]            =						OpenLlamaModel(_UpperCamelCase   )
							scaled_model.to(_UpperCamelCase   )
							scaled_model.eval()
							UpperCAmelCase_    :      Dict            =						scaled_model(_UpperCamelCase   ).last_hidden_state
							UpperCAmelCase_    :      Union[str, 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   )   )














import inspect
import logging
import os
import random
import shutil
import tempfile
import unittest

import pytest
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

from accelerate import Accelerator
from accelerate.test_utils import execute_subprocess_async, require_cuda
from accelerate.utils import ProjectConfiguration, set_seed


__UpperCAmelCase								=						logging.getLogger(__name__)




def      lowercase__      (     __snake_case    :							List[Any]=2     ,       __snake_case    :							Union[str, Any]=3     ,       __snake_case    :							Any=16     ,       __snake_case    :							int = 10     ,       __snake_case    :							int = 2  ):


					'''simple docstring'''


					def get_dataset(__snake_case    :							Optional[Any]  ):
										UpperCAmelCase_    :      Optional[Any]            =						torch.randn(batch_size * n_batches     ,       1  )
										return TensorDataset(__snake_case     ,       a * x + b + 0.1 * torch.randn(batch_size * n_batches     ,       1  )  )

					UpperCAmelCase_    :      Any            =						get_dataset(__snake_case  )
					UpperCAmelCase_    :      str            =						get_dataset(__snake_case  )
					UpperCAmelCase_    :      int            =						DataLoader(__snake_case     ,       shuffle=__snake_case     ,       batch_size=__snake_case     ,       num_workers=4  )
					UpperCAmelCase_    :      int            =						DataLoader(__snake_case     ,       shuffle=__snake_case     ,       batch_size=__snake_case     ,       num_workers=4  )
					return (train_dataloader, valid_dataloader)




def      lowercase__      (     __snake_case    :							Optional[int]     ,       __snake_case    :							str     ,       __snake_case    :							Optional[int]     ,       __snake_case    :							List[str]     ,       __snake_case    :							Any     ,       __snake_case    :							Tuple=None  ):


					'''simple docstring'''


					UpperCAmelCase_    :      Optional[int]            =						[]
					for epoch in range(__snake_case  ):
										# Train quickly
										model.train()
										for batch in dataloader:
															UpperCAmelCase_       ,  UpperCAmelCase_    :      List[Any]            =						batch
															UpperCAmelCase_    :      List[Any]            =						model(__snake_case  )
															UpperCAmelCase_    :      int            =						torch.nn.functional.mse_loss(__snake_case     ,       __snake_case  )
															accelerator.backward(__snake_case  )
															optimizer.step()
															optimizer.zero_grad()
										rands.append(random.random()  )  # Introduce some randomness
										if scheduler is not None:
															scheduler.step()
					return rands



class   lowerCamelCase      (nn.Module			):







		'''simple docstring'''



		def __init__(							self   )	->							Optional[Any]:
							super().__init__()
							UpperCAmelCase_    :      List[Any]            =						nn.Parameter(torch.randn(1   )   )
							UpperCAmelCase_    :      Optional[int]            =						nn.Parameter(torch.randn(1   )   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase   )	->							Optional[Any]:
							return x * self.a + self.b





class   lowerCamelCase      (unittest.TestCase			):







		'''simple docstring'''



		def 					__UpperCAmelCase			(							self   )	->							Dict:
							with tempfile.TemporaryDirectory() as tmpdir:
												set_seed(4_2   )
												UpperCAmelCase_    :      Tuple            =						DummyModel()
												UpperCAmelCase_    :      List[str]            =						torch.optim.Adam(params=model.parameters()				,		lr=1E-3   )
												UpperCAmelCase_       ,  UpperCAmelCase_    :      Union[str, Any]            =						dummy_dataloaders()
												UpperCAmelCase_    :      Optional[int]            =						ProjectConfiguration(total_limit=1				,		project_dir=_UpperCamelCase				,		automatic_checkpoint_naming=_UpperCamelCase   )
												# Train baseline
												UpperCAmelCase_    :      Dict            =						Accelerator(project_config=_UpperCamelCase   )
												UpperCAmelCase_       ,  UpperCAmelCase_       ,  UpperCAmelCase_       ,  UpperCAmelCase_    :      str            =						accelerator.prepare(
												    _UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )
												# Save initial
												accelerator.save_state()

												# Save second state
												accelerator.save_state()
												self.assertEqual(len(os.listdir(accelerator.project_dir   )   )				,		1   )



		def 					__UpperCAmelCase			(							self   )	->							int:
							with tempfile.TemporaryDirectory() as tmpdir:
												set_seed(4_2   )
												UpperCAmelCase_    :      Optional[Any]            =						DummyModel()
												UpperCAmelCase_    :      str            =						torch.optim.Adam(params=model.parameters()				,		lr=1E-3   )
												UpperCAmelCase_       ,  UpperCAmelCase_    :      Tuple            =						dummy_dataloaders()
												# Train baseline
												UpperCAmelCase_    :      Tuple            =						Accelerator()
												UpperCAmelCase_       ,  UpperCAmelCase_       ,  UpperCAmelCase_       ,  UpperCAmelCase_    :      Optional[Any]            =						accelerator.prepare(
												    _UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )
												# Save initial
												UpperCAmelCase_    :      Any            =						os.path.join(_UpperCamelCase				,		'initial'   )
												accelerator.save_state(_UpperCamelCase   )
												((UpperCAmelCase_)       ,  (UpperCAmelCase_))    :      Optional[int]            =						model.a.item(), model.b.item()
												UpperCAmelCase_    :      Dict            =						optimizer.state_dict()
												UpperCAmelCase_    :      Union[str, Any]            =						train(3				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )
												((UpperCAmelCase_)       ,  (UpperCAmelCase_))    :      Union[str, Any]            =						model.a.item(), model.b.item()
												UpperCAmelCase_    :      Any            =						optimizer.state_dict()

												# Train partially
												set_seed(4_2   )
												UpperCAmelCase_    :      int            =						DummyModel()
												UpperCAmelCase_    :      int            =						torch.optim.Adam(params=model.parameters()				,		lr=1E-3   )
												UpperCAmelCase_       ,  UpperCAmelCase_    :      str            =						dummy_dataloaders()
												UpperCAmelCase_    :      Optional[Any]            =						Accelerator()
												UpperCAmelCase_       ,  UpperCAmelCase_       ,  UpperCAmelCase_       ,  UpperCAmelCase_    :      Tuple            =						accelerator.prepare(
												    _UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )
												accelerator.load_state(_UpperCamelCase   )
												((UpperCAmelCase_)       ,  (UpperCAmelCase_))    :      List[str]            =						model.a.item(), model.b.item()
												UpperCAmelCase_    :      Optional[Any]            =						optimizer.state_dict()
												self.assertEqual(_UpperCamelCase				,		_UpperCamelCase   )
												self.assertEqual(_UpperCamelCase				,		_UpperCamelCase   )
												self.assertEqual(_UpperCamelCase				,		_UpperCamelCase   )

												UpperCAmelCase_    :      Dict            =						train(2				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )
												# Save everything
												UpperCAmelCase_    :      Union[str, Any]            =						os.path.join(_UpperCamelCase				,		'checkpoint'   )
												accelerator.save_state(_UpperCamelCase   )

												# Load everything back in and make sure all states work
												accelerator.load_state(_UpperCamelCase   )
												test_rands += train(1				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )
												((UpperCAmelCase_)       ,  (UpperCAmelCase_))    :      Optional[Any]            =						model.a.item(), model.b.item()
												UpperCAmelCase_    :      Union[str, Any]            =						optimizer.state_dict()
												self.assertEqual(_UpperCamelCase				,		_UpperCamelCase   )
												self.assertEqual(_UpperCamelCase				,		_UpperCamelCase   )
												self.assertEqual(_UpperCamelCase				,		_UpperCamelCase   )
												self.assertEqual(_UpperCamelCase				,		_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							int:
							with tempfile.TemporaryDirectory() as tmpdir:
												set_seed(4_2   )
												UpperCAmelCase_    :      Tuple            =						DummyModel()
												UpperCAmelCase_    :      Optional[int]            =						torch.optim.Adam(params=model.parameters()				,		lr=1E-3   )
												UpperCAmelCase_       ,  UpperCAmelCase_    :      Optional[int]            =						dummy_dataloaders()
												UpperCAmelCase_    :      Any            =						ProjectConfiguration(automatic_checkpoint_naming=_UpperCamelCase   )

												# Train baseline
												UpperCAmelCase_    :      str            =						Accelerator(project_dir=_UpperCamelCase				,		project_config=_UpperCamelCase   )
												UpperCAmelCase_       ,  UpperCAmelCase_       ,  UpperCAmelCase_       ,  UpperCAmelCase_    :      Any            =						accelerator.prepare(
												    _UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )
												# Save initial
												accelerator.save_state()
												((UpperCAmelCase_)       ,  (UpperCAmelCase_))    :      Optional[int]            =						model.a.item(), model.b.item()
												UpperCAmelCase_    :      Optional[int]            =						optimizer.state_dict()
												UpperCAmelCase_    :      Optional[Any]            =						train(3				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )
												((UpperCAmelCase_)       ,  (UpperCAmelCase_))    :      Tuple            =						model.a.item(), model.b.item()
												UpperCAmelCase_    :      Optional[int]            =						optimizer.state_dict()

												# Train partially
												set_seed(4_2   )
												UpperCAmelCase_    :      Any            =						DummyModel()
												UpperCAmelCase_    :      Any            =						torch.optim.Adam(params=model.parameters()				,		lr=1E-3   )
												UpperCAmelCase_       ,  UpperCAmelCase_    :      Union[str, Any]            =						dummy_dataloaders()
												UpperCAmelCase_    :      Tuple            =						ProjectConfiguration(iteration=1				,		automatic_checkpoint_naming=_UpperCamelCase   )
												UpperCAmelCase_    :      List[Any]            =						Accelerator(project_dir=_UpperCamelCase				,		project_config=_UpperCamelCase   )
												UpperCAmelCase_       ,  UpperCAmelCase_       ,  UpperCAmelCase_       ,  UpperCAmelCase_    :      Any            =						accelerator.prepare(
												    _UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )
												accelerator.load_state(os.path.join(_UpperCamelCase				,		'checkpoints'				,		'checkpoint_0'   )   )
												((UpperCAmelCase_)       ,  (UpperCAmelCase_))    :      str            =						model.a.item(), model.b.item()
												UpperCAmelCase_    :      List[Any]            =						optimizer.state_dict()
												self.assertEqual(_UpperCamelCase				,		_UpperCamelCase   )
												self.assertEqual(_UpperCamelCase				,		_UpperCamelCase   )
												self.assertEqual(_UpperCamelCase				,		_UpperCamelCase   )

												UpperCAmelCase_    :      Union[str, Any]            =						train(2				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )
												# Save everything
												accelerator.save_state()

												# Load everything back in and make sure all states work
												accelerator.load_state(os.path.join(_UpperCamelCase				,		'checkpoints'				,		'checkpoint_1'   )   )
												test_rands += train(1				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )
												((UpperCAmelCase_)       ,  (UpperCAmelCase_))    :      List[Any]            =						model.a.item(), model.b.item()
												UpperCAmelCase_    :      Dict            =						optimizer.state_dict()
												self.assertEqual(_UpperCamelCase				,		_UpperCamelCase   )
												self.assertEqual(_UpperCamelCase				,		_UpperCamelCase   )
												self.assertEqual(_UpperCamelCase				,		_UpperCamelCase   )
												self.assertEqual(_UpperCamelCase				,		_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							Dict:
							UpperCAmelCase_    :      Optional[Any]            =						torch.tensor([1, 2, 3]   )
							UpperCAmelCase_    :      Any            =						torch.tensor([2, 3, 4]   )
							UpperCAmelCase_    :      Union[str, Any]            =						DummyModel()
							UpperCAmelCase_    :      List[str]            =						torch.optim.Adam(net.parameters()   )
							UpperCAmelCase_    :      Any            =						Accelerator()
							with self.assertRaises(_UpperCamelCase   ) as ve:
												accelerator.register_for_checkpointing(_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )
							UpperCAmelCase_    :      Optional[int]            =						str(ve.exception   )
							self.assertTrue('Item at index 0' in message   )
							self.assertTrue('Item at index 1' in message   )
							self.assertFalse('Item at index 2' in message   )
							self.assertFalse('Item at index 3' in message   )



		def 					__UpperCAmelCase			(							self   )	->							int:
							with tempfile.TemporaryDirectory() as tmpdir:
												set_seed(4_2   )
												UpperCAmelCase_    :      int            =						DummyModel()
												UpperCAmelCase_    :      Any            =						torch.optim.Adam(params=model.parameters()				,		lr=1E-3   )
												UpperCAmelCase_    :      Dict            =						torch.optim.lr_scheduler.StepLR(_UpperCamelCase				,		step_size=1				,		gamma=0.99   )
												UpperCAmelCase_       ,  UpperCAmelCase_    :      Tuple            =						dummy_dataloaders()
												UpperCAmelCase_    :      Tuple            =						ProjectConfiguration(automatic_checkpoint_naming=_UpperCamelCase   )
												# Train baseline
												UpperCAmelCase_    :      Tuple            =						Accelerator(project_dir=_UpperCamelCase				,		project_config=_UpperCamelCase   )
												UpperCAmelCase_       ,  UpperCAmelCase_       ,  UpperCAmelCase_       ,  UpperCAmelCase_       ,  UpperCAmelCase_    :      Any            =						accelerator.prepare(
												    _UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )
												# Save initial
												accelerator.save_state()
												UpperCAmelCase_    :      Dict            =						scheduler.state_dict()
												train(3				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )
												self.assertNotEqual(_UpperCamelCase				,		scheduler.state_dict()   )

												# Load everything back in and make sure all states work
												accelerator.load_state(os.path.join(_UpperCamelCase				,		'checkpoints'				,		'checkpoint_0'   )   )
												self.assertEqual(_UpperCamelCase				,		scheduler.state_dict()   )



		def 					__UpperCAmelCase			(							self   )	->							Dict:
							with tempfile.TemporaryDirectory() as tmpdir:
												set_seed(4_2   )
												UpperCAmelCase_    :      Optional[int]            =						DummyModel()
												UpperCAmelCase_    :      Dict            =						ProjectConfiguration(automatic_checkpoint_naming=_UpperCamelCase				,		total_limit=2   )
												# Train baseline
												UpperCAmelCase_    :      Optional[int]            =						Accelerator(project_dir=_UpperCamelCase				,		project_config=_UpperCamelCase   )
												UpperCAmelCase_    :      str            =						accelerator.prepare(_UpperCamelCase   )
												# Save 3 states:
												for _ in range(1_1   ):
																	accelerator.save_state()
												self.assertTrue(not os.path.exists(os.path.join(_UpperCamelCase				,		'checkpoints'				,		'checkpoint_0'   )   )   )
												self.assertTrue(os.path.exists(os.path.join(_UpperCamelCase				,		'checkpoints'				,		'checkpoint_9'   )   )   )
												self.assertTrue(os.path.exists(os.path.join(_UpperCamelCase				,		'checkpoints'				,		'checkpoint_10'   )   )   )



		@require_cuda
		def 					__UpperCAmelCase			(							self   )	->							str:
							UpperCAmelCase_    :      List[str]            =						['torchrun', f"--nproc_per_node={torch.cuda.device_count()}", inspect.getfile(self.__class__   )]
							execute_subprocess_async(_UpperCamelCase				,		env=os.environ.copy()   )


if __name__ == "__main__":
		__UpperCAmelCase								=						'/tmp/accelerate/state_checkpointing'
		__UpperCAmelCase								=						DummyModel()
		__UpperCAmelCase								=						torch.optim.Adam(params=model.parameters(), lr=1E-3)
		__UpperCAmelCase								=						torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9_9)
		__UpperCAmelCase							,   __UpperCAmelCase								=						dummy_dataloaders()
		__UpperCAmelCase								=						ProjectConfiguration(automatic_checkpoint_naming=True)
		# Train baseline
		__UpperCAmelCase								=						Accelerator(project_dir=savedir, project_config=project_config, mixed_precision='no')
		if accelerator.process_index == 0:
				if os.path.exists(savedir):
						shutil.rmtree(savedir)
				os.makedirs(savedir)
		__UpperCAmelCase							,   __UpperCAmelCase							,   __UpperCAmelCase							,   __UpperCAmelCase							,   __UpperCAmelCase								=						accelerator.prepare(
		    model, optimizer, train_dataloader, valid_dataloader, scheduler
		)
		__UpperCAmelCase							,   __UpperCAmelCase								=						accelerator.prepare(model, optimizer)
		train(3, model, train_dataloader, optimizer, accelerator, scheduler)
		# Check that the intial optimizer is loaded on the GPU
		for group in optimizer.param_groups:
				__UpperCAmelCase								=						group['params'][0].device
				break
		assert param_device.type == accelerator.device.type
		__UpperCAmelCase								=						model.cpu()
		accelerator.wait_for_everyone()
		accelerator.save_state()
		accelerator.wait_for_everyone()

		# Check CPU state
		accelerator.load_state(os.path.join(savedir, 'checkpoints', 'checkpoint_0'), map_location='cpu')
		for group in optimizer.param_groups:
				__UpperCAmelCase								=						group['params'][0].device
				break
		assert (
		    param_device.type == torch.device('cpu').type
		), F"Loaded optimizer states did not match, expected to be loaded on the CPU but got {param_device}"

		# Check device state
		model.to(accelerator.device)
		accelerator.load_state(os.path.join(savedir, 'checkpoints', 'checkpoint_0'), map_location='on_device')
		for group in optimizer.param_groups:
				__UpperCAmelCase								=						group['params'][0].device
				break
		assert (
		    param_device.type == accelerator.device.type
		), F"Loaded optimizer states did not match, expected to be loaded on {accelerator.device} but got {param_device}"

		# Check error
		with pytest.raises(TypeError, match='Unsupported optimizer map location passed'):
				accelerator.load_state(os.path.join(savedir, 'checkpoints', 'checkpoint_0'), map_location='invalid')
		accelerator.wait_for_everyone()
		if accelerator.process_index == 0:
				shutil.rmtree(savedir)
		accelerator.wait_for_everyone()






