






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    =       namedtuple('ExampleNamedTuple', 'a b c')
class        __snake_case(					unittest.TestCase      ):



 def 						UpperCAmelCase    (	self   :   List[str]	):

    """simple docstring"""






    lowercase__						:	Optional[Any]         = torch.randn(5 ,		2	)
    lowercase__						:	str         = torch.device("""cuda"""	) if torch.cuda.is_available() else torch.device("""cpu"""	)

    lowercase__						:	List[Any]         = send_to_device(UpperCAmelCase__ ,		UpperCAmelCase__	)
    self.assertTrue(torch.equal(resulta.cpu() ,		UpperCAmelCase__	)	)

    lowercase__						:	Optional[int]         = 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	)

    lowercase__						:	Any         = 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	)

    lowercase__						:	Tuple         = 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	)

    lowercase__						:	List[str]         = 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 						UpperCAmelCase    (	self   :   Optional[Any]	):

    """simple docstring"""






    with self.assertRaises(UpperCAmelCase__	) as cm:
       lowercase__						:	Tuple         = 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 						UpperCAmelCase    (	self   :   Dict	):

    """simple docstring"""






    lowercase__						:	Any         = torch.tensor([1, 2, 3, 4, 5]	)
    self.assertEqual(listify(UpperCAmelCase__	) ,		[1, 2, 3, 4, 5]	)

    lowercase__						:	List[str]         = 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]]	)

    lowercase__						:	Dict         = 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 						UpperCAmelCase    (	self   :   Optional[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 						UpperCAmelCase    (	self   :   Any	):

    """simple docstring"""






    lowercase__						:	Optional[Any]         = RegressionModel()
    lowercase__						:	Union[str, Any]         = model.forward
    lowercase__						:	List[str]         = convert_outputs_to_fpaa(model.forward	)
    lowercase__						:	Union[str, Any]         = extract_model_from_parallel(UpperCAmelCase__ ,		keep_fpaa_wrapper=UpperCAmelCase__	)
    lowercase__						:	Dict         = pickle.dumps(UpperCAmelCase__	)



 @require_cuda
 def 						UpperCAmelCase    (	self   :   Union[str, Any]	):

    """simple docstring"""






    lowercase__						:	Union[str, Any]         = RegressionModel()
    lowercase__						:	int         = model.forward
    lowercase__						:	List[str]         = torch.cuda.amp.autocast(dtype=torch.floataa	)(model.forward	)
    lowercase__						:	int         = convert_outputs_to_fpaa(model.forward	)
    lowercase__						:	List[str]         = extract_model_from_parallel(UpperCAmelCase__ ,		keep_fpaa_wrapper=UpperCAmelCase__	)
    lowercase__						:	Any         = pickle.dumps(UpperCAmelCase__	)



 @require_cuda
 @require_torch_min_version(version="""2.0"""	)
 def 						UpperCAmelCase    (	self   :   Any	):

    """simple docstring"""






    lowercase__						:	str         = RegressionModel()
    lowercase__						:	List[Any]         = model.forward
    lowercase__						:	Optional[Any]         = torch.cuda.amp.autocast(dtype=torch.floataa	)(model.forward	)
    lowercase__						:	Tuple         = convert_outputs_to_fpaa(model.forward	)
    lowercase__						:	Tuple         = torch.compile(model.forward ,		backend="""inductor"""	)
    lowercase__						:	List[Any]         = torch.randn(4 ,		10	).cuda()
    lowercase__						:	Tuple         = model(UpperCAmelCase__	)



 def 						UpperCAmelCase    (	self   :   Dict	):

    """simple docstring"""






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

    self.assertEqual(UpperCAmelCase__ ,		UpperCAmelCase__	)



 @require_torch_min_version(version="""2.0"""	)
 def 						UpperCAmelCase    (	self   :   Tuple	):

    """simple docstring"""






    lowercase__						:	Union[str, Any]         = RegressionModel()
    lowercase__						:	Tuple         = torch.compile(UpperCAmelCase__	)

    # could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
    lowercase__						:	List[Any]         = torch.nn.parallel.DataParallel(UpperCAmelCase__	)
    lowercase__						:	Optional[int]         = torch.compile(UpperCAmelCase__	)
    lowercase__						:	Any         = extract_model_from_parallel(UpperCAmelCase__	)

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







 def 						UpperCAmelCase    (	self   :   List[Any]	):

    """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"""]	)	)










from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


a    =       {'configuration_nat': ['NAT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'NatConfig']}


try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 a    =       [
     'NAT_PRETRAINED_MODEL_ARCHIVE_LIST',
     'NatForImageClassification',
     'NatModel',
     'NatPreTrainedModel',
     'NatBackbone',
 ]

if TYPE_CHECKING:
 from .configuration_nat import NAT_PRETRAINED_CONFIG_ARCHIVE_MAP, NatConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_nat import (
      NAT_PRETRAINED_MODEL_ARCHIVE_LIST,
      NatBackbone,
      NatForImageClassification,
      NatModel,
      NatPreTrainedModel,
  )

else:
 import sys

 a    =       _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)


