



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


SCREAMING_SNAKE_CASE_:Dict             =						logging.getLogger(__name__)




def 	__UpperCamelCase	(  _lowerCAmelCase=2      ,     _lowerCAmelCase=3      ,     _lowerCAmelCase=16      ,     _lowerCAmelCase = 10      ,     _lowerCAmelCase = 2		)       ->  int:



    """simple docstring"""






    def get_dataset(_lowerCAmelCase		):
        A :       Any          =   torch.randn(batch_size * n_batches      ,     1		)
        return TensorDataset(_lowerCAmelCase      ,     a * x + b + 0.1 * torch.randn(batch_size * n_batches      ,     1		)		)

    A :       List[Any]          =   get_dataset(_lowerCAmelCase		)
    A :       Any          =   get_dataset(_lowerCAmelCase		)
    A :       str          =   DataLoader(_lowerCAmelCase      ,     shuffle=_lowerCAmelCase      ,     batch_size=_lowerCAmelCase      ,     num_workers=4		)
    A :       int          =   DataLoader(_lowerCAmelCase      ,     shuffle=_lowerCAmelCase      ,     batch_size=_lowerCAmelCase      ,     num_workers=4		)
    return (train_dataloader, valid_dataloader)




def 	__UpperCamelCase	(  _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase=None		)       ->  Tuple:



    """simple docstring"""






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






class SCREAMING_SNAKE_CASE__     (		nn.Module    ):




    '''simple docstring'''






    def __init__(       self		):
        super().__init__()
        A :       Tuple          =   nn.Parameter(torch.randn(1		)		)
        A :       Dict          =   nn.Parameter(torch.randn(1		)		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        return x * self.a + self.b






class SCREAMING_SNAKE_CASE__     (		unittest.TestCase    ):




    '''simple docstring'''






    def 			_lowerCAmelCase					(       self		):
        with tempfile.TemporaryDirectory() as tmpdir:
            set_seed(42		)
            A :       Tuple          =   DummyModel()
            A :       Tuple          =   torch.optim.Adam(params=model.parameters(),		lr=1e-3		)
            A     ,							A :       List[Any]          =   dummy_dataloaders()
            A :       str          =   ProjectConfiguration(total_limit=1,		project_dir=lowerCamelCase__,		automatic_checkpoint_naming=lowerCamelCase__		)
            # Train baseline
            A :       List[Any]          =   Accelerator(project_config=lowerCamelCase__		)
            A     ,							A     ,							A     ,							A :       int          =   accelerator.prepare(
                lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		)
            # Save initial
            accelerator.save_state()

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






    def 			_lowerCAmelCase					(       self		):
        with tempfile.TemporaryDirectory() as tmpdir:
            set_seed(42		)
            A :       List[Any]          =   DummyModel()
            A :       List[str]          =   torch.optim.Adam(params=model.parameters(),		lr=1e-3		)
            A     ,							A :       Optional[int]          =   dummy_dataloaders()
            # Train baseline
            A :       Optional[int]          =   Accelerator()
            A     ,							A     ,							A     ,							A :       str          =   accelerator.prepare(
                lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		)
            # Save initial
            A :       int          =   os.path.join(lowerCamelCase__,		"""initial"""		)
            accelerator.save_state(lowerCamelCase__		)
            ((A)     ,							(A)) :       int          =   model.a.item(), model.b.item()
            A :       List[Any]          =   optimizer.state_dict()
            A :       Dict          =   train(3,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		)
            ((A)     ,							(A)) :       Union[str, Any]          =   model.a.item(), model.b.item()
            A :       str          =   optimizer.state_dict()

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

            A :       int          =   train(2,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		)
            # Save everything
            A :       Optional[Any]          =   os.path.join(lowerCamelCase__,		"""checkpoint"""		)
            accelerator.save_state(lowerCamelCase__		)

            # Load everything back in and make sure all states work
            accelerator.load_state(lowerCamelCase__		)
            test_rands += train(1,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		)
            ((A)     ,							(A)) :       int          =   model.a.item(), model.b.item()
            A :       List[Any]          =   optimizer.state_dict()
            self.assertEqual(lowerCamelCase__,		lowerCamelCase__		)
            self.assertEqual(lowerCamelCase__,		lowerCamelCase__		)
            self.assertEqual(lowerCamelCase__,		lowerCamelCase__		)
            self.assertEqual(lowerCamelCase__,		lowerCamelCase__		)






    def 			_lowerCAmelCase					(       self		):
        with tempfile.TemporaryDirectory() as tmpdir:
            set_seed(42		)
            A :       List[Any]          =   DummyModel()
            A :       Optional[Any]          =   torch.optim.Adam(params=model.parameters(),		lr=1e-3		)
            A     ,							A :       Tuple          =   dummy_dataloaders()
            A :       str          =   ProjectConfiguration(automatic_checkpoint_naming=lowerCamelCase__		)

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

            # Train partially
            set_seed(42		)
            A :       List[Any]          =   DummyModel()
            A :       str          =   torch.optim.Adam(params=model.parameters(),		lr=1e-3		)
            A     ,							A :       int          =   dummy_dataloaders()
            A :       Tuple          =   ProjectConfiguration(iteration=1,		automatic_checkpoint_naming=lowerCamelCase__		)
            A :       Any          =   Accelerator(project_dir=lowerCamelCase__,		project_config=lowerCamelCase__		)
            A     ,							A     ,							A     ,							A :       Any          =   accelerator.prepare(
                lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		)
            accelerator.load_state(os.path.join(lowerCamelCase__,		"""checkpoints""",		"""checkpoint_0"""		)		)
            ((A)     ,							(A)) :       List[str]          =   model.a.item(), model.b.item()
            A :       Optional[int]          =   optimizer.state_dict()
            self.assertEqual(lowerCamelCase__,		lowerCamelCase__		)
            self.assertEqual(lowerCamelCase__,		lowerCamelCase__		)
            self.assertEqual(lowerCamelCase__,		lowerCamelCase__		)

            A :       List[Any]          =   train(2,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		)
            # Save everything
            accelerator.save_state()

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






    def 			_lowerCAmelCase					(       self		):
        A :       str          =   torch.tensor([1, 2, 3]		)
        A :       Optional[int]          =   torch.tensor([2, 3, 4]		)
        A :       Optional[int]          =   DummyModel()
        A :       Tuple          =   torch.optim.Adam(net.parameters()		)
        A :       Optional[int]          =   Accelerator()
        with self.assertRaises(lowerCamelCase__		) as ve:
            accelerator.register_for_checkpointing(lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		)
        A :       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 			_lowerCAmelCase					(       self		):
        with tempfile.TemporaryDirectory() as tmpdir:
            set_seed(42		)
            A :       int          =   DummyModel()
            A :       Any          =   torch.optim.Adam(params=model.parameters(),		lr=1e-3		)
            A :       List[Any]          =   torch.optim.lr_scheduler.StepLR(lowerCamelCase__,		step_size=1,		gamma=0.99		)
            A     ,							A :       Tuple          =   dummy_dataloaders()
            A :       Dict          =   ProjectConfiguration(automatic_checkpoint_naming=lowerCamelCase__		)
            # Train baseline
            A :       Tuple          =   Accelerator(project_dir=lowerCamelCase__,		project_config=lowerCamelCase__		)
            A     ,							A     ,							A     ,							A     ,							A :       Any          =   accelerator.prepare(
                lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		)
            # Save initial
            accelerator.save_state()
            A :       int          =   scheduler.state_dict()
            train(3,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		)
            self.assertNotEqual(lowerCamelCase__,		scheduler.state_dict()		)

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






    def 			_lowerCAmelCase					(       self		):
        with tempfile.TemporaryDirectory() as tmpdir:
            set_seed(42		)
            A :       int          =   DummyModel()
            A :       Optional[Any]          =   ProjectConfiguration(automatic_checkpoint_naming=lowerCamelCase__,		total_limit=2		)
            # Train baseline
            A :       str          =   Accelerator(project_dir=lowerCamelCase__,		project_config=lowerCamelCase__		)
            A :       Tuple          =   accelerator.prepare(lowerCamelCase__		)
            # Save 3 states:
            for _ in range(11		):
                accelerator.save_state()
            self.assertTrue(not os.path.exists(os.path.join(lowerCamelCase__,		"""checkpoints""",		"""checkpoint_0"""		)		)		)
            self.assertTrue(os.path.exists(os.path.join(lowerCamelCase__,		"""checkpoints""",		"""checkpoint_9"""		)		)		)
            self.assertTrue(os.path.exists(os.path.join(lowerCamelCase__,		"""checkpoints""",		"""checkpoint_10"""		)		)		)






    @require_cuda
    def 			_lowerCAmelCase					(       self		):
        A :       int          =   ["""torchrun""", f'''--nproc_per_node={torch.cuda.device_count()}''', inspect.getfile(self.__class__		)]
        execute_subprocess_async(lowerCamelCase__,		env=os.environ.copy()		)


if __name__ == "__main__":
     SCREAMING_SNAKE_CASE_:Any             =						"""/tmp/accelerate/state_checkpointing"""
     SCREAMING_SNAKE_CASE_:Tuple             =						DummyModel()
     SCREAMING_SNAKE_CASE_:Tuple             =						torch.optim.Adam(params=model.parameters(), lr=1E-3)
     SCREAMING_SNAKE_CASE_:Dict             =						torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.99)
     SCREAMING_SNAKE_CASE_       ,     SCREAMING_SNAKE_CASE_:Tuple             =						dummy_dataloaders()
     SCREAMING_SNAKE_CASE_:Optional[int]             =						ProjectConfiguration(automatic_checkpoint_naming=True)
     # Train baseline
     SCREAMING_SNAKE_CASE_:Dict             =						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)
     SCREAMING_SNAKE_CASE_       ,     SCREAMING_SNAKE_CASE_       ,     SCREAMING_SNAKE_CASE_       ,     SCREAMING_SNAKE_CASE_       ,     SCREAMING_SNAKE_CASE_:Optional[Any]             =						accelerator.prepare(
         model, optimizer, train_dataloader, valid_dataloader, scheduler
     )
     SCREAMING_SNAKE_CASE_       ,     SCREAMING_SNAKE_CASE_:str             =						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:
          SCREAMING_SNAKE_CASE_:Union[str, Any]             =						group["""params"""][0].device
          break
     assert param_device.type == accelerator.device.type
     SCREAMING_SNAKE_CASE_:Optional[int]             =						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:
          SCREAMING_SNAKE_CASE_:List[str]             =						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:
          SCREAMING_SNAKE_CASE_:str             =						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()







import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
     import jax.numpy as jnp

     from transformers.models.albert.modeling_flax_albert import (
         FlaxAlbertForMaskedLM,
         FlaxAlbertForMultipleChoice,
         FlaxAlbertForPreTraining,
         FlaxAlbertForQuestionAnswering,
         FlaxAlbertForSequenceClassification,
         FlaxAlbertForTokenClassification,
         FlaxAlbertModel,
     )






class SCREAMING_SNAKE_CASE__     (		unittest.TestCase    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__=13,		lowerCamelCase__=7,		lowerCamelCase__=True,		lowerCamelCase__=True,		lowerCamelCase__=True,		lowerCamelCase__=True,		lowerCamelCase__=99,		lowerCamelCase__=32,		lowerCamelCase__=5,		lowerCamelCase__=4,		lowerCamelCase__=37,		lowerCamelCase__="gelu",		lowerCamelCase__=0.1,		lowerCamelCase__=0.1,		lowerCamelCase__=512,		lowerCamelCase__=16,		lowerCamelCase__=2,		lowerCamelCase__=0.02,		lowerCamelCase__=4,		):
        A :       List[str]          =   parent
        A :       Optional[int]          =   batch_size
        A :       Union[str, Any]          =   seq_length
        A :       Any          =   is_training
        A :       List[str]          =   use_attention_mask
        A :       Union[str, Any]          =   use_token_type_ids
        A :       Any          =   use_labels
        A :       str          =   vocab_size
        A :       Union[str, Any]          =   hidden_size
        A :       str          =   num_hidden_layers
        A :       List[Any]          =   num_attention_heads
        A :       Optional[int]          =   intermediate_size
        A :       Optional[Any]          =   hidden_act
        A :       Dict          =   hidden_dropout_prob
        A :       List[Any]          =   attention_probs_dropout_prob
        A :       Optional[int]          =   max_position_embeddings
        A :       int          =   type_vocab_size
        A :       str          =   type_sequence_label_size
        A :       List[Any]          =   initializer_range
        A :       str          =   num_choices






    def 			_lowerCAmelCase					(       self		):
        A :       Optional[int]          =   ids_tensor([self.batch_size, self.seq_length],		self.vocab_size		)

        A :       Union[str, Any]          =   None
        if self.use_attention_mask:
            A :       Optional[int]          =   random_attention_mask([self.batch_size, self.seq_length]		)

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

        A :       Optional[int]          =   AlbertConfig(
            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=lowerCamelCase__,		initializer_range=self.initializer_range,		)

        return config, input_ids, token_type_ids, attention_mask






    def 			_lowerCAmelCase					(       self		):
        A :       Dict          =   self.prepare_config_and_inputs()
        A     ,							A     ,							A     ,							A :       str          =   config_and_inputs
        A :       Union[str, Any]          =   {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
        return config, inputs_dict






@require_flax
class SCREAMING_SNAKE_CASE__     (		SCREAMING_SNAKE_CASE__ ,						unittest.TestCase    ):




    '''simple docstring'''






    __lowerCamelCase			: Any					=      (
        (
            FlaxAlbertModel,
            FlaxAlbertForPreTraining,
            FlaxAlbertForMaskedLM,
            FlaxAlbertForMultipleChoice,
            FlaxAlbertForQuestionAnswering,
            FlaxAlbertForSequenceClassification,
            FlaxAlbertForTokenClassification,
            FlaxAlbertForQuestionAnswering,
        )
        if is_flax_available()
        else ()
    )






    def 			_lowerCAmelCase					(       self		):
        A :       Dict          =   FlaxAlbertModelTester(self		)






    @slow
    def 			_lowerCAmelCase					(       self		):
        for model_class_name in self.all_model_classes:
            A :       Dict          =   model_class_name.from_pretrained("""albert-base-v2"""		)
            A :       Optional[int]          =   model(np.ones((1, 1)		)		)
            self.assertIsNotNone(lowerCamelCase__		)






@require_flax
class SCREAMING_SNAKE_CASE__     (		unittest.TestCase    ):




    '''simple docstring'''






    @slow
    def 			_lowerCAmelCase					(       self		):
        A :       Dict          =   FlaxAlbertModel.from_pretrained("""albert-base-v2"""		)
        A :       List[str]          =   np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]		)
        A :       str          =   np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]		)
        A :       Optional[int]          =   model(lowerCamelCase__,		attention_mask=lowerCamelCase__		)[0]
        A :       str          =   (1, 11, 768)
        self.assertEqual(output.shape,		lowerCamelCase__		)
        A :       Optional[int]          =   np.array(
            [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]]		)

        self.assertTrue(jnp.allclose(output[:, 1:4, 1:4],		lowerCamelCase__,		atol=1e-4		)		)


