
'''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,
)


a__		:   Dict            =       namedtuple("ExampleNamedTuple", "a b c")


class  UpperCamelCase__  (    unittest.TestCase):






 def 	lowercase_						(			self  :str			)				->			List[str]:







       '''simple docstring'''



       __A       	=			torch.randn(5	,				2			)
       __A       	=			torch.device('cuda'			) if torch.cuda.is_available() else torch.device('cpu'			)

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

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

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

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

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






 def 	lowercase_						(			self  :Optional[int]			)				->			List[str]:







       '''simple docstring'''



       with self.assertRaises(_A			) as cm:
             __A       	=			recursively_apply(torch.tensor	,				(torch.tensor(1			), 1)	,				error_on_other_type=_A			)
       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 	lowercase_						(			self  :str			)				->			List[Any]:







       '''simple docstring'''



       __A       	=			torch.tensor([1, 2, 3, 4, 5]			)
       self.assertEqual(listify(_A			)	,				[1, 2, 3, 4, 5]			)

       __A       	=			torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]			)
       self.assertEqual(listify(_A			)	,				[[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(_A			)	,				[[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]]			)






 def 	lowercase_						(			self  :Optional[int]			)				->			List[Any]:







       '''simple docstring'''



       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 	lowercase_						(			self  :int			)				->			Any:







       '''simple docstring'''



       __A       	=			RegressionModel()
       __A       	=			model.forward
       __A       	=			convert_outputs_to_fpaa(model.forward			)
       __A       	=			extract_model_from_parallel(_A	,				keep_fpaa_wrapper=_A			)
       __A       	=			pickle.dumps(_A			)






 @require_cuda
 def 	lowercase_						(			self  :List[str]			)				->			str:







       '''simple docstring'''



       __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(_A	,				keep_fpaa_wrapper=_A			)
       __A       	=			pickle.dumps(_A			)






 @require_cuda
 @require_torch_min_version(version='2.0'			)
 def 	lowercase_						(			self  :Optional[Any]			)				->			Optional[int]:







       '''simple docstring'''



       __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(_A			)






 def 	lowercase_						(			self  :List[Any]			)				->			int:







       '''simple docstring'''



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

       self.assertEqual(_A	,				_A			)






 @require_torch_min_version(version='2.0'			)
 def 	lowercase_						(			self  :List[str]			)				->			Dict:







       '''simple docstring'''



       __A       	=			RegressionModel()
       __A       	=			torch.compile(_A			)

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

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






 def 	lowercase_						(			self  :Dict			)				->			Dict:







       '''simple docstring'''



       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'''





from .hash_table import HashTable
from .number_theory.prime_numbers import is_prime, next_prime


class  UpperCamelCase__  (    SCREAMING_SNAKE_CASE):






 def __init__(			self  :Union[str, Any]	,				*_A  :int	,				**_A  :Optional[Any]			)				->			Union[str, Any]:







       '''simple docstring'''



       super().__init__(*_A	,				**_A			)






 def 	lowercase_						(			self  :Dict	,				_A  :Union[str, Any]	,				_A  :Dict			)				->			Dict:







       '''simple docstring'''



       __A       	=			(
           next_prime(value % self.size_table			)
           if not is_prime(value % self.size_table			)
           else value % self.size_table
       )  # gt = bigger than
       return next_prime_gt - (data % next_prime_gt)






 def 	lowercase_						(			self  :List[Any]	,				_A  :int	,				_A  :Tuple	,				_A  :Union[str, Any]			)				->			List[Any]:







       '''simple docstring'''



       return (increment * self.__hash_function_a(_A	,				_A			)) % self.size_table






 def 	lowercase_						(			self  :List[Any]	,				_A  :List[str]	,				_A  :str=None			)				->			List[Any]:







       '''simple docstring'''



       __A       	=			1
       __A       	=			self.hash_function(_A			)

       while self.values[new_key] is not None and self.values[new_key] != key:
             __A       	=			(
                 self.__hash_double_function(_A	,				_A	,				_A			)
                 if self.balanced_factor() >= self.lim_charge
                 else None
             )
             if new_key is None:
                   break
             else:
                   i += 1

       return new_key






