





'''simple docstring'''






import json
import os
import tempfile
from unittest.mock import patch

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

from accelerate import DistributedType, infer_auto_device_map, init_empty_weights
from accelerate.accelerator import Accelerator
from accelerate.state import GradientState, PartialState
from accelerate.test_utils import require_bnb, require_multi_gpu, slow
from accelerate.test_utils.testing import AccelerateTestCase, require_cuda
from accelerate.utils import patch_environment
def   UpperCamelCase__ (						):





      """simple docstring"""


      _lowerCAmelCase					=				torch.nn.Linear(2				,      4    )
      _lowerCAmelCase					=				torch.optim.AdamW(model.parameters()				,      lr=1.0    )
      _lowerCAmelCase					=				torch.optim.lr_scheduler.OneCycleLR(lowerCAmelCase				,      max_lr=0.01				,      steps_per_epoch=2				,      epochs=1    )
      _lowerCAmelCase					=				DataLoader(TensorDataset(torch.tensor([1, 2, 3]    )    )    )
      _lowerCAmelCase					=				DataLoader(TensorDataset(torch.tensor([4, 5, 6]    )    )    )

      return model, optimizer, scheduler, train_dl, valid_dl
def   UpperCamelCase__ (						lowerCAmelCase    ):





      """simple docstring"""


      return (model.weight.abs().sum() + model.bias.abs().sum()).item()
def   UpperCamelCase__ (						lowerCAmelCase    ):





      """simple docstring"""


      _lowerCAmelCase					=				torch.nn.Linear(*tuple(model.weight.T.shape    )    ).state_dict()
      model.load_state_dict(lowerCAmelCase    )





class 					UpperCAmelCase	(     snake_case_  ):




      @require_cuda
      def 							lowercase__       (     self     :					Tuple    )							-> Dict:
            _lowerCAmelCase					=				Accelerator()
            assert PartialState._shared_state["_cpu"] is False
            assert PartialState._shared_state["device"].type == "cuda"
            with self.assertRaises(__snake_case    ):
                  _lowerCAmelCase					=				Accelerator(cpu=__snake_case    )




      def 							lowercase__       (     self     :					List[str]    )							-> Union[str, Any]:
            _lowerCAmelCase					=				Accelerator()
            _lowerCAmelCase					=				GradientState()
            assert state.num_steps == 1
            _lowerCAmelCase					=				4
            assert state.num_steps == 4

            assert state.sync_gradients is True
            _lowerCAmelCase					=				False
            assert state.sync_gradients is False
            GradientState._reset_state()




      def 							lowercase__       (     self     :					int    )							-> Optional[int]:
            _lowerCAmelCase					=				Accelerator()
            _lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase					=				create_components()

            (
                (
                _lowerCAmelCase
            )	,				(
                _lowerCAmelCase
            )	,				(
                _lowerCAmelCase
            )	,				(
                _lowerCAmelCase
            )	,				(
                _lowerCAmelCase
            )	,				
            )					=				accelerator.prepare(__snake_case     ,     __snake_case     ,     __snake_case     ,     __snake_case     ,     __snake_case    )

            self.assertTrue(prepared_model in accelerator._models    )
            self.assertTrue(prepared_optimizer in accelerator._optimizers    )
            self.assertTrue(prepared_scheduler in accelerator._schedulers    )
            self.assertTrue(prepared_train_dl in accelerator._dataloaders    )
            self.assertTrue(prepared_valid_dl in accelerator._dataloaders    )




      def 							lowercase__       (     self     :					List[Any]    )							-> List[str]:
            _lowerCAmelCase					=				Accelerator()
            _lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase					=				create_components()
            accelerator.prepare(__snake_case     ,     __snake_case     ,     __snake_case     ,     __snake_case     ,     __snake_case    )
            accelerator.free_memory()

            self.assertTrue(len(accelerator._models    ) == 0    )
            self.assertTrue(len(accelerator._optimizers    ) == 0    )
            self.assertTrue(len(accelerator._schedulers    ) == 0    )
            self.assertTrue(len(accelerator._dataloaders    ) == 0    )




      def 							lowercase__       (     self     :					int    )							-> Union[str, Any]:
            PartialState._reset_state()

            # Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
            def noop(*__snake_case     :					Optional[int]     ,     **__snake_case     :					Union[str, Any]    ):
                  pass

            with patch("""torch.cuda.set_device"""     ,     __snake_case    ), patch_environment(ACCELERATE_TORCH_DEVICE="""cuda:64"""    ):
                  _lowerCAmelCase					=				Accelerator()
                  self.assertEqual(str(accelerator.state.device    )     ,     """cuda:64"""    )




      def 							lowercase__       (     self     :					Dict    )							-> str:
            _lowerCAmelCase					=				Accelerator()
            _lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase					=				create_components()
            accelerator.prepare(__snake_case     ,     __snake_case     ,     __snake_case     ,     __snake_case     ,     __snake_case    )

            _lowerCAmelCase					=				get_signature(__snake_case    )

            with tempfile.TemporaryDirectory() as tmpdirname:
                  accelerator.save_state(__snake_case    )

                  # make sure random weights don't match
                  load_random_weights(__snake_case    )
                  self.assertTrue(abs(model_signature - get_signature(__snake_case    )    ) > 1E-3    )

                  # make sure loaded weights match
                  accelerator.load_state(__snake_case    )
                  self.assertTrue(abs(model_signature - get_signature(__snake_case    )    ) < 1E-3    )




      def 							lowercase__       (     self     :					Dict    )							-> List[str]:
            _lowerCAmelCase					=				Accelerator()
            _lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase					=				create_components()
            accelerator.prepare(__snake_case     ,     __snake_case     ,     __snake_case     ,     __snake_case     ,     __snake_case    )

            _lowerCAmelCase					=				get_signature(__snake_case    )

            # saving hook
            def save_config(__snake_case     :					str     ,     __snake_case     :					int     ,     __snake_case     :					Any    ):
                  _lowerCAmelCase					=				{"""class_name""": models[0].__class__.__name__}

                  with open(os.path.join(__snake_case     ,     """data.json"""    )     ,     """w"""    ) as f:
                        json.dump(__snake_case     ,     __snake_case    )

        # loading hook
            def load_config(__snake_case     :					str     ,     __snake_case     :					List[Any]    ):
                  with open(os.path.join(__snake_case     ,     """data.json"""    )     ,     """r"""    ) as f:
                        _lowerCAmelCase					=				json.load(__snake_case    )

                  _lowerCAmelCase					=				config["""class_name"""]

            _lowerCAmelCase					=				accelerator.register_save_state_pre_hook(__snake_case    )
            _lowerCAmelCase					=				accelerator.register_load_state_pre_hook(__snake_case    )

            with tempfile.TemporaryDirectory() as tmpdirname:
                  accelerator.save_state(__snake_case    )

                  # make sure random weights don't match with hooks
                  load_random_weights(__snake_case    )
                  self.assertTrue(abs(model_signature - get_signature(__snake_case    )    ) > 1E-3    )

                  # random class name to verify correct one is loaded
                  _lowerCAmelCase					=				"""random"""

                  # make sure loaded weights match with hooks
                  accelerator.load_state(__snake_case    )
                  self.assertTrue(abs(model_signature - get_signature(__snake_case    )    ) < 1E-3    )

                  # mode.class_name is loaded from config
                  self.assertTrue(model.class_name == model.__class__.__name__    )

            # remove hooks
            save_hook.remove()
            load_hook.remove()

            with tempfile.TemporaryDirectory() as tmpdirname:
                  accelerator.save_state(__snake_case    )

                  # make sure random weights don't match with hooks removed
                  load_random_weights(__snake_case    )
                  self.assertTrue(abs(model_signature - get_signature(__snake_case    )    ) > 1E-3    )

                  # random class name to verify correct one is loaded
                  _lowerCAmelCase					=				"""random"""

                  # make sure loaded weights match with hooks removed
                  accelerator.load_state(__snake_case    )
                  self.assertTrue(abs(model_signature - get_signature(__snake_case    )    ) < 1E-3    )

                  # mode.class_name is NOT loaded from config
                  self.assertTrue(model.class_name != model.__class__.__name__    )




      def 							lowercase__       (     self     :					Dict    )							-> Tuple:
            _lowerCAmelCase					=				Accelerator()
            _lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase					=				create_components()
            _lowerCAmelCase					=				None

            # This should work
            _lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase					=				accelerator.prepare(
                __snake_case     ,     __snake_case     ,     __snake_case     ,     __snake_case     ,     __snake_case     ,     __snake_case    )
            self.assertTrue(dummy_obj is None    )




      def 							lowercase__       (     self     :					Optional[Any]    )							-> List[Any]:
            _lowerCAmelCase					=				Accelerator()
            _lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase					=				create_components()
            _lowerCAmelCase					=				[1, 2, 3]

            # This should work
            _lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase					=				accelerator.prepare(
                __snake_case     ,     __snake_case     ,     __snake_case     ,     __snake_case     ,     __snake_case     ,     __snake_case    )
            self.assertEqual(
                getattr(__snake_case     ,     """_is_accelerate_prepared"""     ,     __snake_case    )     ,     __snake_case     ,     """Dummy object should have `_is_accelerate_prepared` set to `True`"""     ,     )
            self.assertEqual(
                getattr(__snake_case     ,     """_is_accelerate_prepared"""     ,     __snake_case    )     ,     __snake_case     ,     """Model is missing `_is_accelerator_prepared` or is set to `False`"""     ,     )
            self.assertEqual(
                getattr(__snake_case     ,     """_is_accelerate_prepared"""     ,     __snake_case    )     ,     __snake_case     ,     """Optimizer is missing `_is_accelerator_prepared` or is set to `False`"""     ,     )
            self.assertEqual(
                getattr(__snake_case     ,     """_is_accelerate_prepared"""     ,     __snake_case    )     ,     __snake_case     ,     """Scheduler is missing `_is_accelerator_prepared` or is set to `False`"""     ,     )
            self.assertEqual(
                getattr(__snake_case     ,     """_is_accelerate_prepared"""     ,     __snake_case    )     ,     __snake_case     ,     """Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`"""     ,     )
            self.assertEqual(
                getattr(__snake_case     ,     """_is_accelerate_prepared"""     ,     __snake_case    )     ,     __snake_case     ,     """Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`"""     ,     )




      @slow
      @require_bnb
      def 							lowercase__       (     self     :					List[str]    )							-> Optional[int]:
            from transformers import AutoModelForCausalLM

            _lowerCAmelCase					=				AutoModelForCausalLM.from_pretrained(
                """EleutherAI/gpt-neo-125m"""     ,     load_in_abit=__snake_case     ,     device_map={"""""": 0}     ,     )
            _lowerCAmelCase					=				Accelerator()

            # This should work
            _lowerCAmelCase					=				accelerator.prepare(__snake_case    )




      @slow
      @require_bnb
      def 							lowercase__       (     self     :					int    )							-> str:
            from transformers import AutoModelForCausalLM

            _lowerCAmelCase					=				Accelerator()

            with init_empty_weights():
                  _lowerCAmelCase					=				AutoModelForCausalLM.from_pretrained(
                      """EleutherAI/gpt-neo-125m"""     ,     )
                  model.tie_weights()
                  _lowerCAmelCase					=				infer_auto_device_map(__snake_case    )
                  _lowerCAmelCase					=				"""cpu"""

            _lowerCAmelCase					=				AutoModelForCausalLM.from_pretrained(
                """EleutherAI/gpt-neo-125m"""     ,     device_map=__snake_case     ,     load_in_abit=__snake_case     ,     llm_inta_enable_fpaa_cpu_offload=__snake_case    )

            # This should not work and get value error
            with self.assertRaises(__snake_case    ):
                  _lowerCAmelCase					=				accelerator.prepare(__snake_case    )




      @slow
      @require_bnb
      @require_multi_gpu
      def 							lowercase__       (     self     :					str    )							-> int:
            from transformers import AutoModelForCausalLM

            _lowerCAmelCase					=				{"""distributed_type""": DistributedType.MULTI_GPU}

            with init_empty_weights():
                  _lowerCAmelCase					=				AutoModelForCausalLM.from_pretrained(
                      """EleutherAI/gpt-neo-125m"""     ,     )
                  model.tie_weights()
                  _lowerCAmelCase					=				infer_auto_device_map(__snake_case    )
                  _lowerCAmelCase					=				1

            _lowerCAmelCase					=				AutoModelForCausalLM.from_pretrained(
                """EleutherAI/gpt-neo-125m"""     ,     load_in_abit=__snake_case     ,     device_map=__snake_case     ,     )
            _lowerCAmelCase					=				Accelerator()

            # This should not work and get value error
            with self.assertRaises(__snake_case    ):
                  _lowerCAmelCase					=				accelerator.prepare(__snake_case    )

            PartialState._reset_state()




      @slow
      @require_bnb
      @require_multi_gpu
      def 							lowercase__       (     self     :					Optional[int]    )							-> Union[str, Any]:
            from transformers import AutoModelForCausalLM

            with init_empty_weights():
                  _lowerCAmelCase					=				AutoModelForCausalLM.from_pretrained(
                      """EleutherAI/gpt-neo-125m"""     ,     )
                  _lowerCAmelCase					=				infer_auto_device_map(__snake_case    )
                  _lowerCAmelCase					=				1

            _lowerCAmelCase					=				AutoModelForCausalLM.from_pretrained(
                """EleutherAI/gpt-neo-125m"""     ,     load_in_abit=__snake_case     ,     device_map=__snake_case     ,     )
            _lowerCAmelCase					=				Accelerator()

            # This should work
            _lowerCAmelCase					=				accelerator.prepare(__snake_case    )




      @require_cuda
      def 							lowercase__       (     self     :					List[str]    )							-> Tuple:
            _lowerCAmelCase					=				torch.nn.Linear(10     ,     10    )
            _lowerCAmelCase					=				torch.optim.SGD(model.parameters()     ,     lr=0.01    )
            _lowerCAmelCase					=				Accelerator(cpu=__snake_case    )
            _lowerCAmelCase					=				accelerator.prepare(__snake_case    )













'''simple docstring'''






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


A__							:  int    ={
    '''configuration_wav2vec2''': ['''WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''Wav2Vec2Config'''],
    '''feature_extraction_wav2vec2''': ['''Wav2Vec2FeatureExtractor'''],
    '''processing_wav2vec2''': ['''Wav2Vec2Processor'''],
    '''tokenization_wav2vec2''': ['''Wav2Vec2CTCTokenizer''', '''Wav2Vec2Tokenizer'''],
}


try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      A__							:  str    =[
          '''WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST''',
          '''Wav2Vec2ForAudioFrameClassification''',
          '''Wav2Vec2ForCTC''',
          '''Wav2Vec2ForMaskedLM''',
          '''Wav2Vec2ForPreTraining''',
          '''Wav2Vec2ForSequenceClassification''',
          '''Wav2Vec2ForXVector''',
          '''Wav2Vec2Model''',
          '''Wav2Vec2PreTrainedModel''',
      ]

try:
      if not is_tf_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      A__							:  Union[str, Any]    =[
          '''TF_WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST''',
          '''TFWav2Vec2ForCTC''',
          '''TFWav2Vec2Model''',
          '''TFWav2Vec2PreTrainedModel''',
          '''TFWav2Vec2ForSequenceClassification''',
      ]

try:
      if not is_flax_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      A__							:  List[str]    =[
          '''FlaxWav2Vec2ForCTC''',
          '''FlaxWav2Vec2ForPreTraining''',
          '''FlaxWav2Vec2Model''',
          '''FlaxWav2Vec2PreTrainedModel''',
      ]


if TYPE_CHECKING:
      from .configuration_wavaveca import WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP, WavaVecaConfig
      from .feature_extraction_wavaveca import WavaVecaFeatureExtractor
      from .processing_wavaveca import WavaVecaProcessor
      from .tokenization_wavaveca import WavaVecaCTCTokenizer, WavaVecaTokenizer

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_wavaveca import (
                WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST,
                WavaVecaForAudioFrameClassification,
                WavaVecaForCTC,
                WavaVecaForMaskedLM,
                WavaVecaForPreTraining,
                WavaVecaForSequenceClassification,
                WavaVecaForXVector,
                WavaVecaModel,
                WavaVecaPreTrainedModel,
            )

      try:
            if not is_tf_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_tf_wavaveca import (
                TF_WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST,
                TFWavaVecaForCTC,
                TFWavaVecaForSequenceClassification,
                TFWavaVecaModel,
                TFWavaVecaPreTrainedModel,
            )

      try:
            if not is_flax_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_tf_wavaveca import (
                FlaxWavaVecaForCTC,
                FlaxWavaVecaForPreTraining,
                FlaxWavaVecaModel,
                FlaxWavaVecaPreTrainedModel,
            )


else:
      import sys

      A__							:  Any    =_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)






