'''simple docstring'''






import os
import pickle
import unittest
from collections import UserDict, namedtuple

import torch

from accelerate.test_utils.testing import require_cuda, require_torch_min_version
from accelerate.test_utils.training import RegressionModel
from accelerate.utils import (
    convert_outputs_to_fpaa,
    extract_model_from_parallel,
    find_device,
    listify,
    patch_environment,
    recursively_apply,
    send_to_device,
)


_UpperCAmelCase						:    Dict     					=					namedtuple('''ExampleNamedTuple''', '''a b c''')







class lowercase_					(		unittest.TestCase ):







   """simple docstring"""

   def   __UpperCAmelCase      (			self			: Any    )     -> str:
        _A								=     torch.randn(5,   2    )
        _A								=     torch.device('cuda'    ) if torch.cuda.is_available() else torch.device('cpu'    )

        _A								=     send_to_device(_lowerCAmelCase,   _lowerCAmelCase    )
        self.assertTrue(torch.equal(resulta.cpu(),   _lowerCAmelCase    )    )

        _A								=     send_to_device((tensor, [tensor, tensor], 1),   _lowerCAmelCase    )
        self.assertIsInstance(_lowerCAmelCase,   _lowerCAmelCase    )
        self.assertTrue(torch.equal(resulta[0].cpu(),   _lowerCAmelCase    )    )
        self.assertIsInstance(resulta[1],   _lowerCAmelCase    )
        self.assertTrue(torch.equal(resulta[1][0].cpu(),   _lowerCAmelCase    )    )
        self.assertTrue(torch.equal(resulta[1][1].cpu(),   _lowerCAmelCase    )    )
        self.assertEqual(resulta[2],   1    )

        _A								=     send_to_device({'a': tensor, 'b': [tensor, tensor], 'c': 1},   _lowerCAmelCase    )
        self.assertIsInstance(_lowerCAmelCase,   _lowerCAmelCase    )
        self.assertTrue(torch.equal(resulta['a'].cpu(),   _lowerCAmelCase    )    )
        self.assertIsInstance(resulta['b'],   _lowerCAmelCase    )
        self.assertTrue(torch.equal(resulta['b'][0].cpu(),   _lowerCAmelCase    )    )
        self.assertTrue(torch.equal(resulta['b'][1].cpu(),   _lowerCAmelCase    )    )
        self.assertEqual(resulta['c'],   1    )

        _A								=     send_to_device(ExampleNamedTuple(a=_lowerCAmelCase,   b=[tensor, tensor],   c=1    ),   _lowerCAmelCase    )
        self.assertIsInstance(_lowerCAmelCase,   _lowerCAmelCase    )
        self.assertTrue(torch.equal(resulta.a.cpu(),   _lowerCAmelCase    )    )
        self.assertIsInstance(resulta.b,   _lowerCAmelCase    )
        self.assertTrue(torch.equal(resulta.b[0].cpu(),   _lowerCAmelCase    )    )
        self.assertTrue(torch.equal(resulta.b[1].cpu(),   _lowerCAmelCase    )    )
        self.assertEqual(resulta.c,   1    )

        _A								=     send_to_device(UserDict({'a': tensor, 'b': [tensor, tensor], 'c': 1}    ),   _lowerCAmelCase    )
        self.assertIsInstance(_lowerCAmelCase,   _lowerCAmelCase    )
        self.assertTrue(torch.equal(resulta['a'].cpu(),   _lowerCAmelCase    )    )
        self.assertIsInstance(resulta['b'],   _lowerCAmelCase    )
        self.assertTrue(torch.equal(resulta['b'][0].cpu(),   _lowerCAmelCase    )    )
        self.assertTrue(torch.equal(resulta['b'][1].cpu(),   _lowerCAmelCase    )    )
        self.assertEqual(resulta['c'],   1    )
   def   __UpperCAmelCase      (			self			: str    )     -> int:
        with self.assertRaises(_lowerCAmelCase    ) as cm:
             _A								=     recursively_apply(torch.tensor,   (torch.tensor(1    ), 1),   error_on_other_type=_lowerCAmelCase    )
        self.assertEqual(
            str(cm.exception    ),   'Unsupported types (<class \'int\'>) passed to `tensor`. Only nested list/tuple/dicts of objects that are valid for `is_torch_tensor` should be passed.',   )
   def   __UpperCAmelCase      (			self			: int    )     -> List[str]:
        _A								=     torch.tensor([1, 2, 3, 4, 5]    )
        self.assertEqual(listify(_lowerCAmelCase    ),   [1, 2, 3, 4, 5]    )

        _A								=     torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]    )
        self.assertEqual(listify(_lowerCAmelCase    ),   [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]    )

        _A								=     torch.tensor([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]]    )
        self.assertEqual(
            listify(_lowerCAmelCase    ),   [[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]]    )
   def   __UpperCAmelCase      (			self			: Any    )     -> Tuple:
        with patch_environment(aa=1,   BB=2    ):
             self.assertEqual(os.environ.get('AA'    ),   '1'    )
             self.assertEqual(os.environ.get('BB'    ),   '2'    )

        self.assertNotIn('AA',   os.environ    )
        self.assertNotIn('BB',   os.environ    )
   def   __UpperCAmelCase      (			self			: Optional[Any]    )     -> Tuple:
        _A								=     RegressionModel()
        _A								=     model.forward
        _A								=     convert_outputs_to_fpaa(model.forward    )
        _A								=     extract_model_from_parallel(_lowerCAmelCase,   keep_fpaa_wrapper=_lowerCAmelCase    )
        _A								=     pickle.dumps(_lowerCAmelCase    )
   @require_cuda
   def   __UpperCAmelCase      (			self			: Optional[int]    )     -> Any:
        _A								=     RegressionModel()
        _A								=     model.forward
        _A								=     torch.cuda.amp.autocast(dtype=torch.floataa    )(model.forward    )
        _A								=     convert_outputs_to_fpaa(model.forward    )
        _A								=     extract_model_from_parallel(_lowerCAmelCase,   keep_fpaa_wrapper=_lowerCAmelCase    )
        _A								=     pickle.dumps(_lowerCAmelCase    )
   @require_cuda
   @require_torch_min_version(version='2.0'    )
   def   __UpperCAmelCase      (			self			: Dict    )     -> int:
        _A								=     RegressionModel()
        _A								=     model.forward
        _A								=     torch.cuda.amp.autocast(dtype=torch.floataa    )(model.forward    )
        _A								=     convert_outputs_to_fpaa(model.forward    )
        _A								=     torch.compile(model.forward,   backend='inductor'    )
        _A								=     torch.randn(4,   10    ).cuda()
        _A								=     model(_lowerCAmelCase    )
   def   __UpperCAmelCase      (			self			: Optional[int]    )     -> Dict:
        _A								=     RegressionModel()
        # could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
        _A								=     torch.nn.parallel.DataParallel(_lowerCAmelCase    )
        _A								=     extract_model_from_parallel(_lowerCAmelCase    )

        self.assertEqual(_lowerCAmelCase,   _lowerCAmelCase    )
   @require_torch_min_version(version='2.0'    )
   def   __UpperCAmelCase      (			self			: int    )     -> List[Any]:
        _A								=     RegressionModel()
        _A								=     torch.compile(_lowerCAmelCase    )

        # could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
        _A								=     torch.nn.parallel.DataParallel(_lowerCAmelCase    )
        _A								=     torch.compile(_lowerCAmelCase    )
        _A								=     extract_model_from_parallel(_lowerCAmelCase    )

        self.assertEqual(compiled_model._orig_mod,   compiled_model_unwrapped._orig_mod    )

   def   __UpperCAmelCase      (			self			: List[Any]    )     -> Optional[int]:
        self.assertEqual(find_device([1, 'a', torch.tensor([1, 2, 3]    )]    ),   torch.device('cpu'    )    )
        self.assertEqual(find_device({'a': 1, 'b': torch.tensor([1, 2, 3]    )}    ),   torch.device('cpu'    )    )
        self.assertIsNone(find_device([1, 'a']    )    )




'''simple docstring'''






import torch

from accelerate import Accelerator, DistributedType







class lowercase_					:







   """simple docstring"""

   def __enter__(			self			: Optional[int]    )     -> Union[str, Any]:
        if self.enabled:
             _A								=     self.model.no_sync()
             self.model_sync_obj.__enter__()

        return self
   def __exit__(			self			: Dict,   UpperCamelCase__			: Any,   UpperCamelCase__			: int,   UpperCamelCase__			: Any    )     -> str:
        if self.enabled:
             # Average all models on exit
             self._sync_and_avg_model_params()
             self.model_sync_obj.__exit__(UpperCamelCase__,   UpperCamelCase__,   UpperCamelCase__    )
   def __init__(			self			: Optional[int],   UpperCamelCase__			: Accelerator,   UpperCamelCase__			: torch.nn.Module,   UpperCamelCase__			: int,   UpperCamelCase__			: bool = True    )     -> str:
        if accelerator.distributed_type not in [
            DistributedType.NO,
            DistributedType.MULTI_CPU,
            DistributedType.MULTI_GPU,
        ]:
             raise NotImplementedError('LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)'    )
        _A								=     enabled and accelerator.distributed_type != DistributedType.NO
        _A								=     0
        if self.enabled:
             _A								=     accelerator
             _A								=     model
             _A								=     local_sgd_steps
   def   __UpperCAmelCase      (			self			: Any    )     -> int:
        self.num_steps += 1
        if not self.enabled:
             return

        if self.num_steps % self.local_sgd_steps == 0:
             self._sync_and_avg_model_params()

   def   __UpperCAmelCase      (			self			: Dict    )     -> int:
        self.accelerator.wait_for_everyone()
        with self.accelerator.autocast():
             for param in self.model.parameters():
                  _A								=     self.accelerator.reduce(param.data,   reduction='mean'    )



