



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_	       =   namedtuple("ExampleNamedTuple", "a b c")







class 							a_  (		unittest.TestCase				):

						def 	__a       (    self	:Any)		->					Optional[int]:
												UpperCAmelCase_     							=     torch.randn(5   ,						2)
												UpperCAmelCase_     							=     torch.device('''cuda''') if torch.cuda.is_available() else torch.device('''cpu''')

												UpperCAmelCase_     							=     send_to_device(_lowercase   ,						_lowercase)
												self.assertTrue(torch.equal(resulta.cpu()   ,						_lowercase))

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

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

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

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

						def 	__a       (    self	:Union[str, Any])		->					List[str]:
												with self.assertRaises(_lowercase) as cm:
																		UpperCAmelCase_     							=     recursively_apply(torch.tensor   ,						(torch.tensor(1), 1)   ,						error_on_other_type=_lowercase)
												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 	__a       (    self	:Union[str, Any])		->					Any:
												UpperCAmelCase_     							=     torch.tensor([1, 2, 3, 4, 5])
												self.assertEqual(listify(_lowercase)   ,						[1, 2, 3, 4, 5])

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

												UpperCAmelCase_     							=     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(_lowercase)   ,						[[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]])

						def 	__a       (    self	:Union[str, Any])		->					List[str]:
												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 	__a       (    self	:Dict)		->					Dict:
												UpperCAmelCase_     							=     RegressionModel()
												UpperCAmelCase_     							=     model.forward
												UpperCAmelCase_     							=     convert_outputs_to_fpaa(model.forward)
												UpperCAmelCase_     							=     extract_model_from_parallel(_lowercase   ,						keep_fpaa_wrapper=_lowercase)
												UpperCAmelCase_     							=     pickle.dumps(_lowercase)

						@require_cuda
						def 	__a       (    self	:Optional[Any])		->					Union[str, Any]:
												UpperCAmelCase_     							=     RegressionModel()
												UpperCAmelCase_     							=     model.forward
												UpperCAmelCase_     							=     torch.cuda.amp.autocast(dtype=torch.floataa)(model.forward)
												UpperCAmelCase_     							=     convert_outputs_to_fpaa(model.forward)
												UpperCAmelCase_     							=     extract_model_from_parallel(_lowercase   ,						keep_fpaa_wrapper=_lowercase)
												UpperCAmelCase_     							=     pickle.dumps(_lowercase)

						@require_cuda
						@require_torch_min_version(version='''2.0''')
						def 	__a       (    self	:Any)		->					str:
												UpperCAmelCase_     							=     RegressionModel()
												UpperCAmelCase_     							=     model.forward
												UpperCAmelCase_     							=     torch.cuda.amp.autocast(dtype=torch.floataa)(model.forward)
												UpperCAmelCase_     							=     convert_outputs_to_fpaa(model.forward)
												UpperCAmelCase_     							=     torch.compile(model.forward   ,						backend='''inductor''')
												UpperCAmelCase_     							=     torch.randn(4   ,						10).cuda()
												UpperCAmelCase_     							=     model(_lowercase)

						def 	__a       (    self	:Dict)		->					int:
												UpperCAmelCase_     							=     RegressionModel()
												# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
												UpperCAmelCase_     							=     torch.nn.parallel.DataParallel(_lowercase)
												UpperCAmelCase_     							=     extract_model_from_parallel(_lowercase)

												self.assertEqual(_lowercase   ,						_lowercase)

						@require_torch_min_version(version='''2.0''')
						def 	__a       (    self	:Optional[Any])		->					Dict:
												UpperCAmelCase_     							=     RegressionModel()
												UpperCAmelCase_     							=     torch.compile(_lowercase)

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

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




						def 	__a       (    self	:Dict)		->					Union[str, Any]:
												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''']))










from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


UpperCamelCase_	       =   {"configuration_vit_hybrid": ["VIT_HYBRID_PRETRAINED_CONFIG_ARCHIVE_MAP", "ViTHybridConfig"]}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCamelCase_	       =   [
							    "VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST",
							    "ViTHybridForImageClassification",
							    "ViTHybridModel",
							    "ViTHybridPreTrainedModel",
							]

try:
							if not is_vision_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCamelCase_	       =   ["ViTHybridImageProcessor"]


if TYPE_CHECKING:
							from .configuration_vit_hybrid import VIT_HYBRID_PRETRAINED_CONFIG_ARCHIVE_MAP, ViTHybridConfig

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_vit_hybrid import (
														    VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST,
														    ViTHybridForImageClassification,
														    ViTHybridModel,
														    ViTHybridPreTrainedModel,
														)

							try:
														if not is_vision_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .image_processing_vit_hybrid import ViTHybridImageProcessor


else:
							import sys

							UpperCamelCase_	       =   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





