




'''simple docstring'''






from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import SwinConfig
from transformers.testing_utils import require_tf, require_vision, slow, to_atuple
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
				import tensorflow as tf

				from transformers.models.swin.modeling_tf_swin import (
				    TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFSwinForImageClassification,
				    TFSwinForMaskedImageModeling,
				    TFSwinModel,
				)


if is_vision_available():
				from PIL import Image

				from transformers import AutoImageProcessor



class    A						:




				'''simple docstring'''





				def __init__(self					,    _UpperCAmelCase					,    _UpperCAmelCase=1_3					,    _UpperCAmelCase=3_2					,    _UpperCAmelCase=2					,    _UpperCAmelCase=3					,    _UpperCAmelCase=1_6					,    _UpperCAmelCase=[1, 2, 1]					,    _UpperCAmelCase=[2, 2, 4]					,    _UpperCAmelCase=2					,    _UpperCAmelCase=2.0					,    _UpperCAmelCase=True					,    _UpperCAmelCase=0.0					,    _UpperCAmelCase=0.0					,    _UpperCAmelCase=0.1					,    _UpperCAmelCase="gelu"					,    _UpperCAmelCase=False					,    _UpperCAmelCase=True					,    _UpperCAmelCase=0.02					,    _UpperCAmelCase=1E-5					,    _UpperCAmelCase=True					,    _UpperCAmelCase=None					,    _UpperCAmelCase=True					,    _UpperCAmelCase=1_0					,    _UpperCAmelCase=8					,    )      ->			None:
							__UpperCamelCase :     Union[str, Any]										=  parent
							__UpperCamelCase :     Optional[int]										=  batch_size
							__UpperCamelCase :     Union[str, Any]										=  image_size
							__UpperCamelCase :     str										=  patch_size
							__UpperCamelCase :     str										=  num_channels
							__UpperCamelCase :     Tuple										=  embed_dim
							__UpperCamelCase :     Optional[int]										=  depths
							__UpperCamelCase :     Optional[int]										=  num_heads
							__UpperCamelCase :     List[Any]										=  window_size
							__UpperCamelCase :     Tuple										=  mlp_ratio
							__UpperCamelCase :     Dict										=  qkv_bias
							__UpperCamelCase :     int										=  hidden_dropout_prob
							__UpperCamelCase :     Tuple										=  attention_probs_dropout_prob
							__UpperCamelCase :     List[str]										=  drop_path_rate
							__UpperCamelCase :     Union[str, Any]										=  hidden_act
							__UpperCamelCase :     Optional[Any]										=  use_absolute_embeddings
							__UpperCamelCase :     Any										=  patch_norm
							__UpperCamelCase :     Optional[Any]										=  layer_norm_eps
							__UpperCamelCase :     List[Any]										=  initializer_range
							__UpperCamelCase :     List[str]										=  is_training
							__UpperCamelCase :     Dict										=  scope
							__UpperCamelCase :     List[str]										=  use_labels
							__UpperCamelCase :     int										=  type_sequence_label_size
							__UpperCamelCase :     Dict										=  encoder_stride
				def 			a_			(self     )      ->			Union[str, Any]:
							__UpperCamelCase :     Tuple										=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

							__UpperCamelCase :     Optional[Any]										=  None
							if self.use_labels:
										__UpperCamelCase :     List[str]										=  ids_tensor([self.batch_size]					,    self.type_sequence_label_size     )

							__UpperCamelCase :     Tuple										=  self.get_config()

							return config, pixel_values, labels
				def 			a_			(self     )      ->			Any:
							return SwinConfig(
							    image_size=self.image_size					,    patch_size=self.patch_size					,    num_channels=self.num_channels					,    embed_dim=self.embed_dim					,    depths=self.depths					,    num_heads=self.num_heads					,    window_size=self.window_size					,    mlp_ratio=self.mlp_ratio					,    qkv_bias=self.qkv_bias					,    hidden_dropout_prob=self.hidden_dropout_prob					,    attention_probs_dropout_prob=self.attention_probs_dropout_prob					,    drop_path_rate=self.drop_path_rate					,    hidden_act=self.hidden_act					,    use_absolute_embeddings=self.use_absolute_embeddings					,    path_norm=self.patch_norm					,    layer_norm_eps=self.layer_norm_eps					,    initializer_range=self.initializer_range					,    encoder_stride=self.encoder_stride					,    )
				def 			a_			(self					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase     )      ->			Optional[int]:
							__UpperCamelCase :     Any										=  TFSwinModel(config=_UpperCAmelCase     )
							__UpperCamelCase :     Any										=  model(_UpperCAmelCase     )

							__UpperCamelCase :     str										=  ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths     ) - 1))
							__UpperCamelCase :     Any										=  int(config.embed_dim * 2 ** (len(config.depths     ) - 1)     )

							self.parent.assertEqual(result.last_hidden_state.shape					,    (self.batch_size, expected_seq_len, expected_dim)     )
				def 			a_			(self					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase     )      ->			Union[str, Any]:
							__UpperCamelCase :     Dict										=  TFSwinForMaskedImageModeling(config=_UpperCAmelCase     )
							__UpperCamelCase :     Union[str, Any]										=  model(_UpperCAmelCase     )
							self.parent.assertEqual(
							    result.logits.shape					,    (self.batch_size, self.num_channels, self.image_size, self.image_size)     )

							# test greyscale images
							__UpperCamelCase :     List[Any]										=  1
							__UpperCamelCase :     Dict										=  TFSwinForMaskedImageModeling(_UpperCAmelCase     )

							__UpperCamelCase :     Union[str, Any]										=  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
							__UpperCamelCase :     str										=  model(_UpperCAmelCase     )
							self.parent.assertEqual(result.logits.shape					,    (self.batch_size, 1, self.image_size, self.image_size)     )
				def 			a_			(self					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase     )      ->			Optional[Any]:
							__UpperCamelCase :     Dict										=  self.type_sequence_label_size
							__UpperCamelCase :     int										=  TFSwinForImageClassification(_UpperCAmelCase     )
							__UpperCamelCase :     Dict										=  model(_UpperCAmelCase					,    labels=_UpperCAmelCase     )
							self.parent.assertEqual(result.logits.shape					,    (self.batch_size, self.type_sequence_label_size)     )

							# test greyscale images
							__UpperCamelCase :     Any										=  1
							__UpperCamelCase :     List[str]										=  TFSwinForImageClassification(_UpperCAmelCase     )
							__UpperCamelCase :     Dict										=  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
							__UpperCamelCase :     Optional[Any]										=  model(_UpperCAmelCase     )
							self.parent.assertEqual(result.logits.shape					,    (self.batch_size, self.type_sequence_label_size)     )





				def 			a_			(self     )      ->			int:
							__UpperCamelCase :     Optional[Any]										=  self.prepare_config_and_inputs()
							__UpperCamelCase						,   __UpperCamelCase						,   __UpperCamelCase :     int										=  config_and_inputs
							__UpperCamelCase :     int										=  {"pixel_values": pixel_values}
							return config, inputs_dict







@require_tf
class    A						(					SCREAMING_SNAKE_CASE__     ,				SCREAMING_SNAKE_CASE__     ,				unittest.TestCase						):




				'''simple docstring'''





				A       =							(
				    (
				        TFSwinModel,
				        TFSwinForImageClassification,
				        TFSwinForMaskedImageModeling,
				    )
				    if is_tf_available()
				    else ()
				)
				A       =							(
				    {"feature-extraction": TFSwinModel, "image-classification": TFSwinForImageClassification}
				    if is_tf_available()
				    else {}
				)

				A       =							False
				A       =							False
				A       =							False
				A       =							False
				def 			a_			(self     )      ->			Dict:
							__UpperCamelCase :     Optional[int]										=  TFSwinModelTester(self     )
							__UpperCamelCase :     Union[str, Any]										=  ConfigTester(self					,    config_class=_UpperCAmelCase					,    embed_dim=3_7     )
				def 			a_			(self     )      ->			Dict:
							self.create_and_test_config_common_properties()
							self.config_tester.create_and_test_config_to_json_string()
							self.config_tester.create_and_test_config_to_json_file()
							self.config_tester.create_and_test_config_from_and_save_pretrained()
							self.config_tester.create_and_test_config_with_num_labels()
							self.config_tester.check_config_can_be_init_without_params()
							self.config_tester.check_config_arguments_init()
				def 			a_			(self     )      ->			Any:
							return
				def 			a_			(self     )      ->			Optional[Any]:
							__UpperCamelCase :     Optional[Any]										=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*_UpperCAmelCase     )
				def 			a_			(self     )      ->			List[Any]:
							__UpperCamelCase :     Any										=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_image_modeling(*_UpperCAmelCase     )
				def 			a_			(self     )      ->			Optional[Any]:
							__UpperCamelCase :     Tuple										=  self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*_UpperCAmelCase     )
				@unittest.skip(reason="Swin does not use inputs_embeds"     )
				def 			a_			(self     )      ->			Dict:
							pass
				def 			a_			(self     )      ->			Optional[int]:
							__UpperCamelCase						,   __UpperCamelCase :     str										=  self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
										__UpperCamelCase :     List[str]										=  model_class(_UpperCAmelCase     )
										self.assertIsInstance(model.get_input_embeddings()					,    tf.keras.layers.Layer     )
										__UpperCamelCase :     Any										=  model.get_output_embeddings()
										self.assertTrue(x is None or isinstance(_UpperCAmelCase					,    tf.keras.layers.Dense     )     )
				def 			a_			(self     )      ->			Dict:
							__UpperCamelCase						,   __UpperCamelCase :     Optional[Any]										=  self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
										__UpperCamelCase :     str										=  model_class(_UpperCAmelCase     )
										__UpperCamelCase :     Any										=  inspect.signature(model.call     )
										# signature.parameters is an OrderedDict => so arg_names order is deterministic
										__UpperCamelCase :     Dict										=  [*signature.parameters.keys()]

										__UpperCamelCase :     str										=  ["pixel_values"]
										self.assertListEqual(arg_names[:1]					,    _UpperCAmelCase     )
				def 			a_			(self     )      ->			Any:
							__UpperCamelCase						,   __UpperCamelCase :     Tuple										=  self.model_tester.prepare_config_and_inputs_for_common()
							__UpperCamelCase :     Any										=  True

							for model_class in self.all_model_classes:
										__UpperCamelCase :     Optional[Any]										=  True
										__UpperCamelCase :     List[Any]										=  False
										__UpperCamelCase :     Optional[Any]										=  True
										__UpperCamelCase :     Union[str, Any]										=  model_class(_UpperCAmelCase     )
										__UpperCamelCase :     Optional[Any]										=  model(**self._prepare_for_class(_UpperCAmelCase					,    _UpperCAmelCase     )     )
										__UpperCamelCase :     str										=  outputs.attentions
										__UpperCamelCase :     Optional[int]										=  len(self.model_tester.depths     )
										self.assertEqual(len(_UpperCAmelCase     )					,    _UpperCAmelCase     )

										# check that output_attentions also work using config
										del inputs_dict["output_attentions"]
										__UpperCamelCase :     Dict										=  True
										__UpperCamelCase :     Optional[int]										=  config.window_size**2
										__UpperCamelCase :     List[Any]										=  model_class(_UpperCAmelCase     )
										__UpperCamelCase :     Union[str, Any]										=  model(**self._prepare_for_class(_UpperCAmelCase					,    _UpperCAmelCase     )     )
										__UpperCamelCase :     Dict										=  outputs.attentions
										self.assertEqual(len(_UpperCAmelCase     )					,    _UpperCAmelCase     )

										self.assertListEqual(
										    list(attentions[0].shape[-3:]     )					,    [self.model_tester.num_heads[0], window_size_squared, window_size_squared]					,    )
										__UpperCamelCase :     Dict										=  len(_UpperCAmelCase     )

										# Check attention is always last and order is fine
										__UpperCamelCase :     Dict										=  True
										__UpperCamelCase :     Optional[Any]										=  True
										__UpperCamelCase :     Union[str, Any]										=  model_class(_UpperCAmelCase     )
										__UpperCamelCase :     List[str]										=  model(**self._prepare_for_class(_UpperCAmelCase					,    _UpperCAmelCase     )     )

										if hasattr(self.model_tester					,    "num_hidden_states_types"     ):
													__UpperCamelCase :     Any										=  self.model_tester.num_hidden_states_types
										else:
													# also another +1 for reshaped_hidden_states
													__UpperCamelCase :     List[Any]										=  2
										self.assertEqual(out_len + added_hidden_states					,    len(_UpperCAmelCase     )     )

										__UpperCamelCase :     Optional[int]										=  outputs.attentions

										self.assertEqual(len(_UpperCAmelCase     )					,    _UpperCAmelCase     )

										self.assertListEqual(
										    list(self_attentions[0].shape[-3:]     )					,    [self.model_tester.num_heads[0], window_size_squared, window_size_squared]					,    )
				def 			a_			(self					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase     )      ->			Dict:
							__UpperCamelCase :     List[str]										=  model_class(_UpperCAmelCase     )
							__UpperCamelCase :     Union[str, Any]										=  model(**self._prepare_for_class(_UpperCAmelCase					,    _UpperCAmelCase     )     )
							__UpperCamelCase :     Any										=  outputs.hidden_states

							__UpperCamelCase :     Union[str, Any]										=  getattr(
							    self.model_tester					,    "expected_num_hidden_layers"					,    len(self.model_tester.depths     ) + 1     )
							self.assertEqual(len(_UpperCAmelCase     )					,    _UpperCAmelCase     )

							# Swin has a different seq_length
							__UpperCamelCase :     List[str]										=  to_atuple(config.patch_size     )

							__UpperCamelCase :     List[str]										=  (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

							self.assertListEqual(
							    list(hidden_states[0].shape[-2:]     )					,    [num_patches, self.model_tester.embed_dim]					,    )

							__UpperCamelCase :     Optional[Any]										=  outputs.reshaped_hidden_states
							self.assertEqual(len(_UpperCAmelCase     )					,    _UpperCAmelCase     )

							__UpperCamelCase						,   __UpperCamelCase						,   __UpperCamelCase						,   __UpperCamelCase :     Any										=  reshaped_hidden_states[0].shape

							__UpperCamelCase :     Tuple										=  tf.reshape(reshaped_hidden_states[0]					,    (batch_size, num_channels, height * width)     )
							__UpperCamelCase :     Optional[int]										=  tf.transpose(_UpperCAmelCase					,    (0, 2, 1)     )

							self.assertListEqual(
							    list(reshaped_hidden_states.shape[-2:]     )					,    [num_patches, self.model_tester.embed_dim]					,    )
				def 			a_			(self     )      ->			List[str]:
							__UpperCamelCase						,   __UpperCamelCase :     Optional[Any]										=  self.model_tester.prepare_config_and_inputs_for_common()

							__UpperCamelCase :     Union[str, Any]										=  to_atuple(self.model_tester.image_size     )

							for model_class in self.all_model_classes:
										__UpperCamelCase :     Tuple										=  True
										self.check_hidden_states_output(_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase     )

										# check that output_hidden_states also work using config
										del inputs_dict["output_hidden_states"]
										__UpperCamelCase :     List[str]										=  True

										self.check_hidden_states_output(_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase     )
				def 			a_			(self     )      ->			List[str]:
							__UpperCamelCase						,   __UpperCamelCase :     Union[str, Any]										=  self.model_tester.prepare_config_and_inputs_for_common()
							__UpperCamelCase :     Optional[Any]										=  3

							__UpperCamelCase :     int										=  to_atuple(self.model_tester.image_size     )
							__UpperCamelCase :     Optional[Any]										=  to_atuple(config.patch_size     )

							__UpperCamelCase :     Dict										=  image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
							__UpperCamelCase :     str										=  image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

							for model_class in self.all_model_classes:
										__UpperCamelCase :     List[Any]										=  True
										self.check_hidden_states_output(_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    (padded_height, padded_width)     )

										# check that output_hidden_states also work using config
										del inputs_dict["output_hidden_states"]
										__UpperCamelCase :     List[str]										=  True
										self.check_hidden_states_output(_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    (padded_height, padded_width)     )





				@slow
				def 			a_			(self     )      ->			str:
							for model_name in TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										__UpperCamelCase :     str										=  TFSwinModel.from_pretrained(_UpperCAmelCase     )
										self.assertIsNotNone(_UpperCAmelCase     )







@require_vision
@require_tf
class    A						(					unittest.TestCase						):




				'''simple docstring'''





				@cached_property
				def 			a_			(self     )      ->			Any:
							return (
							    AutoImageProcessor.from_pretrained("microsoft/swin-tiny-patch4-window7-224"     )
							    if is_vision_available()
							    else None
							)





				@slow
				def 			a_			(self     )      ->			Union[str, Any]:
							__UpperCamelCase :     Dict										=  TFSwinForImageClassification.from_pretrained("microsoft/swin-tiny-patch4-window7-224"     )
							__UpperCamelCase :     Dict										=  self.default_image_processor

							__UpperCamelCase :     List[Any]										=  Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"     )
							__UpperCamelCase :     Dict										=  image_processor(images=_UpperCAmelCase					,    return_tensors="tf"     )

							# forward pass
							__UpperCamelCase :     Optional[int]										=  model(_UpperCAmelCase     )

							# verify the logits
							__UpperCamelCase :     Union[str, Any]										=  tf.TensorShape((1, 1_0_0_0)     )
							self.assertEqual(outputs.logits.shape					,    _UpperCAmelCase     )
							__UpperCamelCase :     List[Any]										=  tf.constant([-0.0_948, -0.6_454, -0.0_921]     )
							self.assertTrue(np.allclose(outputs.logits[0, :3]					,    _UpperCAmelCase					,    atol=1E-4     )     )









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



class    A						(					unittest.TestCase						):




				'''simple docstring'''





				def 			a_			(self     )      ->			Tuple:
							__UpperCamelCase :     Optional[int]										=  torch.randn(5					,    2     )
							__UpperCamelCase :     List[str]										=  torch.device("cuda"     ) if torch.cuda.is_available() else torch.device("cpu"     )

							__UpperCamelCase :     Optional[Any]										=  send_to_device(_UpperCAmelCase					,    _UpperCAmelCase     )
							self.assertTrue(torch.equal(resulta.cpu()					,    _UpperCAmelCase     )     )

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

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

							__UpperCamelCase :     Optional[Any]										=  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     )

							__UpperCamelCase :     Dict										=  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 			a_			(self     )      ->			Optional[Any]:
							with self.assertRaises(_UpperCAmelCase     ) as cm:
										__UpperCamelCase :     int										=  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 			a_			(self     )      ->			Union[str, Any]:
							__UpperCamelCase :     Dict										=  torch.tensor([1, 2, 3, 4, 5]     )
							self.assertEqual(listify(_UpperCAmelCase     )					,    [1, 2, 3, 4, 5]     )

							__UpperCamelCase :     int										=  torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]]     )
							self.assertEqual(listify(_UpperCAmelCase     )					,    [[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]]     )

							__UpperCamelCase :     Dict										=  torch.tensor([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]], [[1_1, 1_2, 1_3, 1_4, 1_5], [1_6, 1_7, 1_8, 1_9, 2_0]]]     )
							self.assertEqual(
							    listify(_UpperCAmelCase     )					,    [[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]], [[1_1, 1_2, 1_3, 1_4, 1_5], [1_6, 1_7, 1_8, 1_9, 2_0]]]     )
				def 			a_			(self     )      ->			Optional[Any]:
							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:
							__UpperCamelCase :     Tuple										=  RegressionModel()
							__UpperCamelCase :     List[str]										=  model.forward
							__UpperCamelCase :     Any										=  convert_outputs_to_fpaa(model.forward     )
							__UpperCamelCase :     List[str]										=  extract_model_from_parallel(_UpperCAmelCase					,    keep_fpaa_wrapper=_UpperCAmelCase     )
							__UpperCamelCase :     Dict										=  pickle.dumps(_UpperCAmelCase     )
				@require_cuda
				def 			a_			(self     )      ->			str:
							__UpperCamelCase :     List[str]										=  RegressionModel()
							__UpperCamelCase :     List[Any]										=  model.forward
							__UpperCamelCase :     List[str]										=  torch.cuda.amp.autocast(dtype=torch.floataa     )(model.forward     )
							__UpperCamelCase :     str										=  convert_outputs_to_fpaa(model.forward     )
							__UpperCamelCase :     int										=  extract_model_from_parallel(_UpperCAmelCase					,    keep_fpaa_wrapper=_UpperCAmelCase     )
							__UpperCamelCase :     int										=  pickle.dumps(_UpperCAmelCase     )
				@require_cuda
				@require_torch_min_version(version="2.0"     )
				def 			a_			(self     )      ->			Tuple:
							__UpperCamelCase :     List[Any]										=  RegressionModel()
							__UpperCamelCase :     Tuple										=  model.forward
							__UpperCamelCase :     Optional[int]										=  torch.cuda.amp.autocast(dtype=torch.floataa     )(model.forward     )
							__UpperCamelCase :     List[Any]										=  convert_outputs_to_fpaa(model.forward     )
							__UpperCamelCase :     Union[str, Any]										=  torch.compile(model.forward					,    backend="inductor"     )
							__UpperCamelCase :     int										=  torch.randn(4					,    1_0     ).cuda()
							__UpperCamelCase :     int										=  model(_UpperCAmelCase     )
				def 			a_			(self     )      ->			Union[str, Any]:
							__UpperCamelCase :     List[Any]										=  RegressionModel()
							# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
							__UpperCamelCase :     Optional[Any]										=  torch.nn.parallel.DataParallel(_UpperCAmelCase     )
							__UpperCamelCase :     Union[str, Any]										=  extract_model_from_parallel(_UpperCAmelCase     )

							self.assertEqual(_UpperCAmelCase					,    _UpperCAmelCase     )
				@require_torch_min_version(version="2.0"     )
				def 			a_			(self     )      ->			Union[str, Any]:
							__UpperCamelCase :     int										=  RegressionModel()
							__UpperCamelCase :     List[str]										=  torch.compile(_UpperCAmelCase     )

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

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





				def 			a_			(self     )      ->			str:
							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"]     )     )



