




import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import MBartConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow
from transformers.utils import cached_property

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
      import os

      # The slow tests are often failing with OOM error on GPU
      # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
      # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
      __SCREAMING_SNAKE_CASE     : int      ="""platform"""

      import jax
      import jax.numpy as jnp

      from transformers import AutoTokenizer
      from transformers.models.mbart.modeling_flax_mbart import (
          FlaxMBartForConditionalGeneration,
          FlaxMBartForQuestionAnswering,
          FlaxMBartForSequenceClassification,
          FlaxMBartModel,
          shift_tokens_right,
      )

def       lowerCAmelCase_      (	lowerCamelCase:							Optional[int]  ,					lowerCamelCase:							Optional[int]  ,					lowerCamelCase:							Optional[int]=None  ,					lowerCamelCase:							str=None  ,					lowerCamelCase:							Any=None  ,					lowerCamelCase:							List[Any]=None  ,					lowerCamelCase:							Tuple=None  ,					lowerCamelCase:							Optional[Any]=None  ,					)       ->      List[str]:


      '''simple docstring'''




      if attention_mask is None:
            lowerCamelCase__   	=					np.where(input_ids != config.pad_token_id  ,					1  ,					0    )
      if decoder_attention_mask is None:
            lowerCamelCase__   	=					np.where(decoder_input_ids != config.pad_token_id  ,					1  ,					0    )
      if head_mask is None:
            lowerCamelCase__   	=					np.ones((config.encoder_layers, config.encoder_attention_heads)    )
      if decoder_head_mask is None:
            lowerCamelCase__   	=					np.ones((config.decoder_layers, config.decoder_attention_heads)    )
      if cross_attn_head_mask is None:
            lowerCamelCase__   	=					np.ones((config.decoder_layers, config.decoder_attention_heads)    )
      return {
          "input_ids": input_ids,
          "decoder_input_ids": decoder_input_ids,
          "attention_mask": attention_mask,
          "decoder_attention_mask": decoder_attention_mask,
      }






class       A      (		unittest.TestCase		):
     """simple docstring"""

     def __init__(     self	: Union[str, Any]						,	A__	: Dict						,	A__	: List[Any]=1_3						,	A__	: Union[str, Any]=7						,	A__	: str=True						,	A__	: Any=False						,	A__	: Tuple=9_9						,	A__	: List[str]=1_6						,	A__	: Optional[int]=2						,	A__	: List[str]=4						,	A__	: List[str]=4						,	A__	: List[Any]="gelu"						,	A__	: Optional[Any]=0.1						,	A__	: Any=0.1						,	A__	: List[Any]=3_2						,	A__	: Optional[int]=2						,	A__	: Optional[Any]=1						,	A__	: Optional[int]=0						,	A__	: str=2						,	A__	: List[str]=0.02						,	)					->							List[Any]:
           lowerCamelCase__   	=					parent
           lowerCamelCase__   	=					batch_size
           lowerCamelCase__   	=					seq_length
           lowerCamelCase__   	=					is_training
           lowerCamelCase__   	=					use_labels
           lowerCamelCase__   	=					vocab_size
           lowerCamelCase__   	=					hidden_size
           lowerCamelCase__   	=					num_hidden_layers
           lowerCamelCase__   	=					num_attention_heads
           lowerCamelCase__   	=					intermediate_size
           lowerCamelCase__   	=					hidden_act
           lowerCamelCase__   	=					hidden_dropout_prob
           lowerCamelCase__   	=					attention_probs_dropout_prob
           lowerCamelCase__   	=					max_position_embeddings
           lowerCamelCase__   	=					eos_token_id
           lowerCamelCase__   	=					pad_token_id
           lowerCamelCase__   	=					bos_token_id
           lowerCamelCase__   	=					decoder_start_token_id
           lowerCamelCase__   	=					initializer_range




     def 			__snake_case (     self	: str      )					->							Union[str, Any]:
           lowerCamelCase__   	=					np.clip(ids_tensor([self.batch_size, self.seq_length - 1]						,	self.vocab_size      )						,	3						,	self.vocab_size      )
           lowerCamelCase__   	=					np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)						,	dtype=np.intaa      ))						,	-1      )

           lowerCamelCase__   	=					shift_tokens_right(UpperCAmelCase_						,	1      )

           lowerCamelCase__   	=					MBartConfig(
               vocab_size=self.vocab_size						,	d_model=self.hidden_size						,	encoder_layers=self.num_hidden_layers						,	decoder_layers=self.num_hidden_layers						,	encoder_attention_heads=self.num_attention_heads						,	decoder_attention_heads=self.num_attention_heads						,	encoder_ffn_dim=self.intermediate_size						,	decoder_ffn_dim=self.intermediate_size						,	dropout=self.hidden_dropout_prob						,	attention_dropout=self.attention_probs_dropout_prob						,	max_position_embeddings=self.max_position_embeddings						,	eos_token_id=self.eos_token_id						,	bos_token_id=self.bos_token_id						,	pad_token_id=self.pad_token_id						,	decoder_start_token_id=self.decoder_start_token_id						,	initializer_range=self.initializer_range						,	use_cache=UpperCAmelCase_						,	)
           lowerCamelCase__   	=					prepare_mbart_inputs_dict(UpperCAmelCase_						,	UpperCAmelCase_						,	UpperCAmelCase_      )
           return config, inputs_dict




     def 			__snake_case (     self	: Dict      )					->							int:
           lowerCamelCase__   	=					self.prepare_config_and_inputs()
           return config, inputs_dict




     def 			__snake_case (     self	: Any						,	A__	: Tuple						,	A__	: List[str]						,	A__	: int      )					->							Union[str, Any]:
           lowerCamelCase__   	=					2_0
           lowerCamelCase__   	=					model_class_name(UpperCAmelCase_      )

           lowerCamelCase__   	=					model.encode(inputs_dict['input_ids']      )

           lowerCamelCase__   	=					(
               inputs_dict["""decoder_input_ids"""],
               inputs_dict["""decoder_attention_mask"""],
           )

           lowerCamelCase__   	=					model.init_cache(decoder_input_ids.shape[0]						,	UpperCAmelCase_						,	UpperCAmelCase_      )
           lowerCamelCase__   	=					jnp.ones((decoder_input_ids.shape[0], max_decoder_length)						,	dtype='i4'      )

           lowerCamelCase__   	=					jnp.broadcast_to(
               jnp.arange(decoder_input_ids.shape[-1] - 1      )[None, :]						,	(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)						,	)
           lowerCamelCase__   	=					model.decode(
               decoder_input_ids[:, :-1]						,	UpperCAmelCase_						,	decoder_attention_mask=UpperCAmelCase_						,	past_key_values=UpperCAmelCase_						,	decoder_position_ids=UpperCAmelCase_						,	)

           lowerCamelCase__   	=					jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]						,	dtype='i4'      )
           lowerCamelCase__   	=					model.decode(
               decoder_input_ids[:, -1:]						,	UpperCAmelCase_						,	decoder_attention_mask=UpperCAmelCase_						,	past_key_values=outputs_cache.past_key_values						,	decoder_position_ids=UpperCAmelCase_						,	)

           lowerCamelCase__   	=					model.decode(UpperCAmelCase_						,	UpperCAmelCase_      )

           lowerCamelCase__   	=					np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])      )      )
           self.parent.assertTrue(diff < 1E-3						,	msg=F"""Max diff is {diff}"""      )




     def 			__snake_case (     self	: int						,	A__	: int						,	A__	: Union[str, Any]						,	A__	: Optional[int]      )					->							Optional[Any]:
           lowerCamelCase__   	=					2_0
           lowerCamelCase__   	=					model_class_name(UpperCAmelCase_      )

           lowerCamelCase__   	=					model.encode(inputs_dict['input_ids']      )

           lowerCamelCase__   	=					(
               inputs_dict["""decoder_input_ids"""],
               inputs_dict["""decoder_attention_mask"""],
           )

           lowerCamelCase__   	=					jnp.concatenate(
               [
                   decoder_attention_mask,
                   jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])      ),
               ]						,	axis=-1						,	)

           lowerCamelCase__   	=					model.init_cache(decoder_input_ids.shape[0]						,	UpperCAmelCase_						,	UpperCAmelCase_      )
           lowerCamelCase__   	=					jnp.broadcast_to(
               jnp.arange(decoder_input_ids.shape[-1] - 1      )[None, :]						,	(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)						,	)

           lowerCamelCase__   	=					model.decode(
               decoder_input_ids[:, :-1]						,	UpperCAmelCase_						,	decoder_attention_mask=UpperCAmelCase_						,	past_key_values=UpperCAmelCase_						,	decoder_position_ids=UpperCAmelCase_						,	)
           lowerCamelCase__   	=					jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]						,	dtype='i4'      )
           lowerCamelCase__   	=					model.decode(
               decoder_input_ids[:, -1:]						,	UpperCAmelCase_						,	past_key_values=outputs_cache.past_key_values						,	decoder_attention_mask=UpperCAmelCase_						,	decoder_position_ids=UpperCAmelCase_						,	)

           lowerCamelCase__   	=					model.decode(UpperCAmelCase_						,	UpperCAmelCase_						,	decoder_attention_mask=UpperCAmelCase_      )

           lowerCamelCase__   	=					np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])      )      )
           self.parent.assertTrue(diff < 1E-3						,	msg=F"""Max diff is {diff}"""      )






@require_flax
class       A      (		unittest.TestCase		):
     """simple docstring"""

     _snake_case:						List[Any]  =99




     def 			__snake_case (     self	: Union[str, Any]      )					->							Dict:
           lowerCamelCase__   	=					np.array(
               [
                   [7_1, 8_2, 1_8, 3_3, 4_6, 9_1, 2],
                   [6_8, 3_4, 2_6, 5_8, 3_0, 8_2, 2],
                   [5, 9_7, 1_7, 3_9, 9_4, 4_0, 2],
                   [7_6, 8_3, 9_4, 2_5, 7_0, 7_8, 2],
                   [8_7, 5_9, 4_1, 3_5, 4_8, 6_6, 2],
                   [5_5, 1_3, 1_6, 5_8, 5, 2, 1],  # note padding
                   [6_4, 2_7, 3_1, 5_1, 1_2, 7_5, 2],
                   [5_2, 6_4, 8_6, 1_7, 8_3, 3_9, 2],
                   [4_8, 6_1, 9, 2_4, 7_1, 8_2, 2],
                   [2_6, 1, 6_0, 4_8, 2_2, 1_3, 2],
                   [2_1, 5, 6_2, 2_8, 1_4, 7_6, 2],
                   [4_5, 9_8, 3_7, 8_6, 5_9, 4_8, 2],
                   [7_0, 7_0, 5_0, 9, 2_8, 0, 2],
               ]						,	dtype=np.intaa						,	)

           lowerCamelCase__   	=					input_ids.shape[0]
           lowerCamelCase__   	=					MBartConfig(
               vocab_size=self.vocab_size						,	d_model=2_4						,	encoder_layers=2						,	decoder_layers=2						,	encoder_attention_heads=2						,	decoder_attention_heads=2						,	encoder_ffn_dim=3_2						,	decoder_ffn_dim=3_2						,	max_position_embeddings=4_8						,	eos_token_id=2						,	pad_token_id=1						,	bos_token_id=0						,	)
           return config, input_ids, batch_size




     def 			__snake_case (     self	: List[Any]      )					->							Tuple:
           lowerCamelCase__   	=					self._get_config_and_data()
           lowerCamelCase__   	=					FlaxMBartForSequenceClassification(UpperCAmelCase_      )
           lowerCamelCase__   	=					model(input_ids=UpperCAmelCase_						,	decoder_input_ids=UpperCAmelCase_      )
           lowerCamelCase__   	=					(batch_size, config.num_labels)
           self.assertEqual(outputs['logits'].shape						,	UpperCAmelCase_      )




     def 			__snake_case (     self	: int      )					->							str:
           lowerCamelCase__   	=					self._get_config_and_data()
           lowerCamelCase__   	=					FlaxMBartForQuestionAnswering(UpperCAmelCase_      )
           lowerCamelCase__   	=					model(input_ids=UpperCAmelCase_      )

           self.assertEqual(outputs['start_logits'].shape						,	input_ids.shape      )
           self.assertEqual(outputs['end_logits'].shape						,	input_ids.shape      )




     def 			__snake_case (     self	: Any      )					->							Union[str, Any]:
           lowerCamelCase__   	=					self._get_config_and_data()
           lowerCamelCase__   	=					FlaxMBartForConditionalGeneration(UpperCAmelCase_      )
           lowerCamelCase__   	=					lm_model(input_ids=UpperCAmelCase_      )
           lowerCamelCase__   	=					(batch_size, input_ids.shape[1], config.vocab_size)
           self.assertEqual(outputs['logits'].shape						,	UpperCAmelCase_      )




     def 			__snake_case (     self	: Dict      )					->							Tuple:
           lowerCamelCase__   	=					MBartConfig(
               vocab_size=self.vocab_size						,	d_model=1_4						,	encoder_layers=2						,	decoder_layers=2						,	encoder_attention_heads=2						,	decoder_attention_heads=2						,	encoder_ffn_dim=8						,	decoder_ffn_dim=8						,	max_position_embeddings=4_8						,	)
           lowerCamelCase__   	=					FlaxMBartForConditionalGeneration(UpperCAmelCase_      )
           lowerCamelCase__   	=					np.array([[7_1, 8_2, 1_8, 3_3, 4_6, 9_1, 2], [6_8, 3_4, 2_6, 5_8, 3_0, 2, 1]]						,	dtype=np.intaa      )
           lowerCamelCase__   	=					np.array([[8_2, 7_1, 8_2, 1_8, 2], [5_8, 6_8, 2, 1, 1]]						,	dtype=np.intaa      )
           lowerCamelCase__   	=					lm_model(input_ids=UpperCAmelCase_						,	decoder_input_ids=UpperCAmelCase_      )
           lowerCamelCase__   	=					(*summary.shape, config.vocab_size)
           self.assertEqual(outputs['logits'].shape						,	UpperCAmelCase_      )




     def 			__snake_case (     self	: Optional[int]      )					->							Union[str, Any]:
           lowerCamelCase__   	=					np.array([[7_1, 8_2, 1_8, 3_3, 2, 1, 1], [6_8, 3_4, 2_6, 5_8, 3_0, 8_2, 2]]						,	dtype=np.intaa      )
           lowerCamelCase__   	=					shift_tokens_right(UpperCAmelCase_						,	1      )
           lowerCamelCase__   	=					np.equal(UpperCAmelCase_						,	1      ).astype(np.floataa      ).sum()
           lowerCamelCase__   	=					np.equal(UpperCAmelCase_						,	1      ).astype(np.floataa      ).sum()
           self.assertEqual(shifted.shape						,	input_ids.shape      )
           self.assertEqual(UpperCAmelCase_						,	n_pad_before - 1      )
           self.assertTrue(np.equal(shifted[:, 0]						,	2      ).all()      )






@require_flax
class       A      (		_lowerCAmelCase      ,    unittest.TestCase      ,    _lowerCAmelCase		):
     """simple docstring"""

     _snake_case:						Dict  =True
     _snake_case:						int  =(
         (
             FlaxMBartModel,
             FlaxMBartForConditionalGeneration,
             FlaxMBartForSequenceClassification,
             FlaxMBartForQuestionAnswering,
         )
         if is_flax_available()
         else ()
     )
     _snake_case:						Dict  =(FlaxMBartForConditionalGeneration,) if is_flax_available() else ()




     def 			__snake_case (     self	: Optional[Any]      )					->							Dict:
           lowerCamelCase__   	=					FlaxMBartModelTester(self      )




     def 			__snake_case (     self	: Dict      )					->							Tuple:
           lowerCamelCase__   	=					self.model_tester.prepare_config_and_inputs()
           for model_class in self.all_model_classes:
                 self.model_tester.check_use_cache_forward(UpperCAmelCase_						,	UpperCAmelCase_						,	UpperCAmelCase_      )




     def 			__snake_case (     self	: List[Any]      )					->							Any:
           lowerCamelCase__   	=					self.model_tester.prepare_config_and_inputs()
           for model_class in self.all_model_classes:
                 self.model_tester.check_use_cache_forward_with_attn_mask(UpperCAmelCase_						,	UpperCAmelCase_						,	UpperCAmelCase_      )




     def 			__snake_case (     self	: str      )					->							Dict:
           lowerCamelCase__   	=					self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                 with self.subTest(model_class.__name__      ):
                       lowerCamelCase__   	=					self._prepare_for_class(UpperCAmelCase_						,	UpperCAmelCase_      )
                       lowerCamelCase__   	=					model_class(UpperCAmelCase_      )

                       @jax.jit
                       def encode_jitted(A__	: Union[str, Any]						,	A__	: List[str]=None						,	**A__	: Union[str, Any]      ):
                             return model.encode(input_ids=UpperCAmelCase_						,	attention_mask=UpperCAmelCase_      )

                       with self.subTest('JIT Enabled'      ):
                             lowerCamelCase__   	=					encode_jitted(**UpperCAmelCase_      ).to_tuple()

                       with self.subTest('JIT Disabled'      ):
                             with jax.disable_jit():
                                   lowerCamelCase__   	=					encode_jitted(**UpperCAmelCase_      ).to_tuple()

                       self.assertEqual(len(UpperCAmelCase_      )						,	len(UpperCAmelCase_      )      )
                       for jitted_output, output in zip(UpperCAmelCase_						,	UpperCAmelCase_      ):
                             self.assertEqual(jitted_output.shape						,	output.shape      )




     def 			__snake_case (     self	: int      )					->							Optional[Any]:
           lowerCamelCase__   	=					self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                 with self.subTest(model_class.__name__      ):
                       lowerCamelCase__   	=					model_class(UpperCAmelCase_      )
                       lowerCamelCase__   	=					model.encode(inputs_dict['input_ids']						,	inputs_dict['attention_mask']      )

                       lowerCamelCase__   	=					{
                           """decoder_input_ids""": inputs_dict["""decoder_input_ids"""],
                           """decoder_attention_mask""": inputs_dict["""decoder_attention_mask"""],
                           """encoder_outputs""": encoder_outputs,
                       }

                       @jax.jit
                       def decode_jitted(A__	: Dict						,	A__	: Tuple						,	A__	: str      ):
                             return model.decode(
                                 decoder_input_ids=UpperCAmelCase_						,	decoder_attention_mask=UpperCAmelCase_						,	encoder_outputs=UpperCAmelCase_						,	)

                       with self.subTest('JIT Enabled'      ):
                             lowerCamelCase__   	=					decode_jitted(**UpperCAmelCase_      ).to_tuple()

                       with self.subTest('JIT Disabled'      ):
                             with jax.disable_jit():
                                   lowerCamelCase__   	=					decode_jitted(**UpperCAmelCase_      ).to_tuple()

                       self.assertEqual(len(UpperCAmelCase_      )						,	len(UpperCAmelCase_      )      )
                       for jitted_output, output in zip(UpperCAmelCase_						,	UpperCAmelCase_      ):
                             self.assertEqual(jitted_output.shape						,	output.shape      )




     @slow
     def 			__snake_case (     self	: List[Any]      )					->							Optional[Any]:
           for model_class_name in self.all_model_classes:
                 lowerCamelCase__   	=					model_class_name.from_pretrained('facebook/mbart-large-cc25'						,	from_pt=UpperCAmelCase_      )
                 # FlaxMBartForSequenceClassification expects eos token in input_ids
                 lowerCamelCase__   	=					np.ones((1, 1)      ) * model.config.eos_token_id
                 lowerCamelCase__   	=					model(UpperCAmelCase_      )
                 self.assertIsNotNone(UpperCAmelCase_      )






@require_flax
@require_sentencepiece
@require_tokenizers
class       A      (		unittest.TestCase		):
     """simple docstring"""

     _snake_case:						Union[str, Any]  =[
         " UN Chief Says There Is No Military Solution in Syria",
     ]
     _snake_case:						List[Any]  =[
         "Şeful ONU declară că nu există o soluţie militară în Siria",
     ]
     _snake_case:						List[Any]  ="facebook/mbart-large-en-ro"




     @cached_property
     def 			__snake_case (     self	: int      )					->							Optional[int]:
           return AutoTokenizer.from_pretrained(self.model_name      )




     @cached_property
     def 			__snake_case (     self	: List[str]      )					->							List[Any]:
           lowerCamelCase__   	=					FlaxMBartForConditionalGeneration.from_pretrained(self.model_name						,	from_pt=UpperCAmelCase_      )
           return model




     def 			__snake_case (     self	: int						,	**A__	: Any      )					->							Any:
           lowerCamelCase__   	=					self.translate_src_text(**UpperCAmelCase_      )
           self.assertListEqual(self.expected_text						,	UpperCAmelCase_      )




     def 			__snake_case (     self	: int						,	**A__	: Dict      )					->							Tuple:
           lowerCamelCase__   	=					self.tokenizer(self.src_text						,	**UpperCAmelCase_						,	return_tensors='np'      )
           lowerCamelCase__   	=					self.model.generate(
               model_inputs.input_ids						,	attention_mask=model_inputs.attention_mask						,	decoder_start_token_id=self.tokenizer.lang_code_to_id['ro_RO']						,	early_stopping=UpperCAmelCase_						,	num_beams=2						,	).sequences
           lowerCamelCase__   	=					self.tokenizer.batch_decode(UpperCAmelCase_						,	skip_special_tokens=UpperCAmelCase_      )
           return generated_words




     @slow
     def 			__snake_case (     self	: Any      )					->							Optional[int]:
           self._assert_generated_batch_equal_expected()




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


UpperCAmelCase:  Union[str, Any]			     =				namedtuple("""ExampleNamedTuple""", """a b c""")





class   UpperCamelCase	(     unittest.TestCase       ):




	"""simple docstring"""






	def    lowerCamelCase__					(     self   ):
			_lowercase				: List[Any]     =   torch.randn(5				,2   )
			_lowercase				: List[Any]     =   torch.device("""cuda"""   ) if torch.cuda.is_available() else torch.device("""cpu"""   )

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

			_lowercase				: 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				: Optional[int]     =   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				: 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    lowerCamelCase__					(     self   ):
			with self.assertRaises(UpperCAmelCase_   ) as cm:
					_lowercase				: str     =   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    lowerCamelCase__					(     self   ):
			_lowercase				: Optional[int]     =   torch.tensor([1, 2, 3, 4, 5]   )
			self.assertEqual(listify(UpperCAmelCase_   )				,[1, 2, 3, 4, 5]   )

			_lowercase				: Optional[int]     =   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				: Optional[Any]     =   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    lowerCamelCase__					(     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    lowerCamelCase__					(     self   ):
			_lowercase				: Optional[Any]     =   RegressionModel()
			_lowercase				: List[Any]     =   model.forward
			_lowercase				: Optional[Any]     =   convert_outputs_to_fpaa(model.forward   )
			_lowercase				: Any     =   extract_model_from_parallel(UpperCAmelCase_				,keep_fpaa_wrapper=UpperCAmelCase_   )
			_lowercase				: Tuple     =   pickle.dumps(UpperCAmelCase_   )






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






	@require_cuda
	@require_torch_min_version(version="""2.0"""   )
	def    lowerCamelCase__					(     self   ):
			_lowercase				: str     =   RegressionModel()
			_lowercase				: Any     =   model.forward
			_lowercase				: Union[str, Any]     =   torch.cuda.amp.autocast(dtype=torch.floataa   )(model.forward   )
			_lowercase				: Optional[Any]     =   convert_outputs_to_fpaa(model.forward   )
			_lowercase				: Union[str, Any]     =   torch.compile(model.forward				,backend="""inductor"""   )
			_lowercase				: Any     =   torch.randn(4				,10   ).cuda()
			_lowercase				: int     =   model(UpperCAmelCase_   )






	def    lowerCamelCase__					(     self   ):
			_lowercase				: List[str]     =   RegressionModel()
			# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
			_lowercase				: Any     =   torch.nn.parallel.DataParallel(UpperCAmelCase_   )
			_lowercase				: List[Any]     =   extract_model_from_parallel(UpperCAmelCase_   )

			self.assertEqual(UpperCAmelCase_				,UpperCAmelCase_   )






	@require_torch_min_version(version="""2.0"""   )
	def    lowerCamelCase__					(     self   ):
			_lowercase				: int     =   RegressionModel()
			_lowercase				: List[str]     =   torch.compile(UpperCAmelCase_   )

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

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






	def    lowerCamelCase__					(     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"""]   )   )



