




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


_lowerCamelCase							:						str								=  namedtuple('ExampleNamedTuple', 'a b c')





class      lowerCamelCase_							(       unittest.TestCase							):




					'''simple docstring'''






					def 				__SCREAMING_SNAKE_CASE	(    self ):
							_snake_case	       =       torch.randn(5,  2 )
							_snake_case	       =       torch.device("cuda" ) if torch.cuda.is_available() else torch.device("cpu" )

							_snake_case	       =       send_to_device(__UpperCamelCase,  __UpperCamelCase )
							self.assertTrue(torch.equal(resulta.cpu(),  __UpperCamelCase ) )

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

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

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

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






					def 				__SCREAMING_SNAKE_CASE	(    self ):
							with self.assertRaises(__UpperCamelCase ) as cm:
									_snake_case	       =       recursively_apply(torch.tensor,  (torch.tensor(1 ), 1),  error_on_other_type=__UpperCamelCase )
							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 				__SCREAMING_SNAKE_CASE	(    self ):
							_snake_case	       =       torch.tensor([1, 2, 3, 4, 5] )
							self.assertEqual(listify(__UpperCamelCase ),  [1, 2, 3, 4, 5] )

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

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






					def 				__SCREAMING_SNAKE_CASE	(    self ):
							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 				__SCREAMING_SNAKE_CASE	(    self ):
							_snake_case	       =       RegressionModel()
							_snake_case	       =       model.forward
							_snake_case	       =       convert_outputs_to_fpaa(model.forward )
							_snake_case	       =       extract_model_from_parallel(__UpperCamelCase,  keep_fpaa_wrapper=__UpperCamelCase )
							_snake_case	       =       pickle.dumps(__UpperCamelCase )






					@require_cuda
					def 				__SCREAMING_SNAKE_CASE	(    self ):
							_snake_case	       =       RegressionModel()
							_snake_case	       =       model.forward
							_snake_case	       =       torch.cuda.amp.autocast(dtype=torch.floataa )(model.forward )
							_snake_case	       =       convert_outputs_to_fpaa(model.forward )
							_snake_case	       =       extract_model_from_parallel(__UpperCamelCase,  keep_fpaa_wrapper=__UpperCamelCase )
							_snake_case	       =       pickle.dumps(__UpperCamelCase )






					@require_cuda
					@require_torch_min_version(version="2.0" )
					def 				__SCREAMING_SNAKE_CASE	(    self ):
							_snake_case	       =       RegressionModel()
							_snake_case	       =       model.forward
							_snake_case	       =       torch.cuda.amp.autocast(dtype=torch.floataa )(model.forward )
							_snake_case	       =       convert_outputs_to_fpaa(model.forward )
							_snake_case	       =       torch.compile(model.forward,  backend="inductor" )
							_snake_case	       =       torch.randn(4,  10 ).cuda()
							_snake_case	       =       model(__UpperCamelCase )






					def 				__SCREAMING_SNAKE_CASE	(    self ):
							_snake_case	       =       RegressionModel()
							# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
							_snake_case	       =       torch.nn.parallel.DataParallel(__UpperCamelCase )
							_snake_case	       =       extract_model_from_parallel(__UpperCamelCase )

							self.assertEqual(__UpperCamelCase,  __UpperCamelCase )






					@require_torch_min_version(version="2.0" )
					def 				__SCREAMING_SNAKE_CASE	(    self ):
							_snake_case	       =       RegressionModel()
							_snake_case	       =       torch.compile(__UpperCamelCase )

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

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






					def 				__SCREAMING_SNAKE_CASE	(    self ):
							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 typing import TYPE_CHECKING

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


_lowerCamelCase							:						Optional[Any]								=  {'configuration_bit': ['BIT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'BitConfig', 'BitOnnxConfig']}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCamelCase							:						Dict								=  [
						    'BIT_PRETRAINED_MODEL_ARCHIVE_LIST',
						    'BitForImageClassification',
						    'BitModel',
						    'BitPreTrainedModel',
						    'BitBackbone',
						]


try:
						if not is_vision_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCamelCase							:						Union[str, Any]								=  ['BitImageProcessor']


if TYPE_CHECKING:
						from .configuration_bit import BIT_PRETRAINED_CONFIG_ARCHIVE_MAP, BitConfig, BitOnnxConfig

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_bit import (
												    BIT_PRETRAINED_MODEL_ARCHIVE_LIST,
												    BitBackbone,
												    BitForImageClassification,
												    BitModel,
												    BitPreTrainedModel,
												)

						try:
												if not is_vision_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .image_processing_bit import BitImageProcessor

else:
						import sys

						_lowerCamelCase							:						Optional[Any]								=  _LazyModule(__name__, globals()['__file__'], _import_structure)





