


import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotConfig, is_flax_available
from transformers.testing_utils import jax_device, require_flax, slow

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
     UpperCamelCase								=					'platform'

     import jax
     import jax.numpy as jnp

     from transformers import BlenderbotTokenizer
     from transformers.models.blenderbot.modeling_flax_blenderbot import (
         FlaxBlenderbotForConditionalGeneration,
         FlaxBlenderbotModel,
         shift_tokens_right,
     )



def 							_A       (						lowerCAmelCase_					: List[Any]       ,    lowerCAmelCase_					: str       ,    lowerCAmelCase_					: str=None       ,    lowerCAmelCase_					: Dict=None       ,    lowerCAmelCase_					: List[str]=None       ,    lowerCAmelCase_					: Optional[Any]=None       ,    lowerCAmelCase_					: int=None       ,    lowerCAmelCase_					: Tuple=None       ,    ):
     """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": attention_mask,
     }





class 			__lowerCamelCase	:




    """simple docstring"""

    def __init__(	self					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]=13       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]=7       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]=True       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]=False       ,     SCREAMING_SNAKE_CASE__					:					Any=99       ,     SCREAMING_SNAKE_CASE__					:					Dict=16       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]=2       ,     SCREAMING_SNAKE_CASE__					:					List[Any]=4       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]=4       ,     SCREAMING_SNAKE_CASE__					:					Dict="gelu"       ,     SCREAMING_SNAKE_CASE__					:					Tuple=0.1       ,     SCREAMING_SNAKE_CASE__					:					int=0.1       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]=32       ,     SCREAMING_SNAKE_CASE__					:					str=2       ,     SCREAMING_SNAKE_CASE__					:					str=1       ,     SCREAMING_SNAKE_CASE__					:					List[str]=0       ,     SCREAMING_SNAKE_CASE__					:					Tuple=0.02       ,     )				->	List[str]:
         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__          =			initializer_range


    def        a				(	self					:					Optional[Any]       )				->	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(_A       ,     1       ,     2       )

         lowerCAmelCase__          =			BlenderbotConfig(
             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       ,     initializer_range=self.initializer_range       ,     use_cache=_A       ,     )
         lowerCAmelCase__          =			prepare_blenderbot_inputs_dict(_A       ,     _A       ,     _A       )
         return config, inputs_dict


    def        a				(	self					:					Any       )				->	int:
         lowerCAmelCase__          =			self.prepare_config_and_inputs()
         return config, inputs_dict


    def        a				(	self					:					int       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					List[Any]       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]       )				->	Optional[int]:
         lowerCAmelCase__          =			20
         lowerCAmelCase__          =			model_class_name(_A       )

         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]       ,     _A       ,     _A       )
         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]       ,     _A       ,     decoder_attention_mask=_A       ,     past_key_values=_A       ,     decoder_position_ids=_A       ,     )

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

         lowerCAmelCase__          =			model.decode(_A       ,     _A       )

         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        a				(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					Any       ,     SCREAMING_SNAKE_CASE__					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					Any       )				->	Tuple:
         lowerCAmelCase__          =			20
         lowerCAmelCase__          =			model_class_name(_A       )

         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]       ,     _A       ,     _A       )
         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]       ,     _A       ,     decoder_attention_mask=_A       ,     past_key_values=_A       ,     decoder_position_ids=_A       ,     )
         lowerCAmelCase__          =			jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]       ,     dtype="i4"       )
         lowerCAmelCase__          =			model.decode(
             decoder_input_ids[:, -1:]       ,     _A       ,     past_key_values=outputs_cache.past_key_values       ,     decoder_attention_mask=_A       ,     decoder_position_ids=_A       ,     )

         lowerCAmelCase__          =			model.decode(_A       ,     _A       ,     decoder_attention_mask=_A       )

         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 			__lowerCamelCase	(   unittest.TestCase   ):




    """simple docstring"""

    snake_case__ =				9_9


    def        a				(	self					:					int       )				->	Optional[int]:
         lowerCAmelCase__          =			np.array(
             [
                 [71, 82, 18, 33, 46, 91, 2],
                 [68, 34, 26, 58, 30, 82, 2],
                 [5, 97, 17, 39, 94, 40, 2],
                 [76, 83, 94, 25, 70, 78, 2],
                 [87, 59, 41, 35, 48, 66, 2],
                 [55, 13, 16, 58, 5, 2, 1],  # note padding
                 [64, 27, 31, 51, 12, 75, 2],
                 [52, 64, 86, 17, 83, 39, 2],
                 [48, 61, 9, 24, 71, 82, 2],
                 [26, 1, 60, 48, 22, 13, 2],
                 [21, 5, 62, 28, 14, 76, 2],
                 [45, 98, 37, 86, 59, 48, 2],
                 [70, 70, 50, 9, 28, 0, 2],
             ]       ,     dtype=np.intaa       ,     )

         lowerCAmelCase__          =			input_ids.shape[0]
         lowerCAmelCase__          =			BlenderbotConfig(
             vocab_size=self.vocab_size       ,     d_model=24       ,     encoder_layers=2       ,     decoder_layers=2       ,     encoder_attention_heads=2       ,     decoder_attention_heads=2       ,     encoder_ffn_dim=32       ,     decoder_ffn_dim=32       ,     max_position_embeddings=48       ,     eos_token_id=2       ,     pad_token_id=1       ,     bos_token_id=0       ,     )
         return config, input_ids, batch_size


    def        a				(	self					:					Tuple       )				->	List[str]:
         lowerCAmelCase__          =			self._get_config_and_data()
         lowerCAmelCase__          =			FlaxBlenderbotForConditionalGeneration(_A       )
         lowerCAmelCase__          =			lm_model(input_ids=_A       )
         lowerCAmelCase__          =			(batch_size, input_ids.shape[1], config.vocab_size)
         self.assertEqual(outputs["logits"].shape       ,     _A       )


    def        a				(	self					:					str       )				->	Union[str, Any]:
         lowerCAmelCase__          =			BlenderbotConfig(
             vocab_size=self.vocab_size       ,     d_model=14       ,     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=48       ,     )
         lowerCAmelCase__          =			FlaxBlenderbotForConditionalGeneration(_A       )
         lowerCAmelCase__          =			np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]       ,     dtype=np.intaa       )
         lowerCAmelCase__          =			np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]       ,     dtype=np.intaa       )
         lowerCAmelCase__          =			lm_model(input_ids=_A       ,     decoder_input_ids=_A       )
         lowerCAmelCase__          =			(*summary.shape, config.vocab_size)
         self.assertEqual(outputs["logits"].shape       ,     _A       )


    def        a				(	self					:					List[str]       )				->	Dict:
         lowerCAmelCase__          =			np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]       ,     dtype=np.intaa       )
         lowerCAmelCase__          =			shift_tokens_right(_A       ,     1       ,     2       )
         lowerCAmelCase__          =			np.equal(_A       ,     1       ).astype(np.floataa       ).sum()
         lowerCAmelCase__          =			np.equal(_A       ,     1       ).astype(np.floataa       ).sum()
         self.assertEqual(shifted.shape       ,     input_ids.shape       )
         self.assertEqual(_A       ,     n_pad_before - 1       )
         self.assertTrue(np.equal(shifted[:, 0]       ,     2       ).all()       )






@require_flax
class 			__lowerCamelCase	(   snake_case_     ,    unittest.TestCase     ,    snake_case_   ):




    """simple docstring"""

    snake_case__ =				True
    snake_case__ =				(
        (
            FlaxBlenderbotModel,
            FlaxBlenderbotForConditionalGeneration,
        )
        if is_flax_available()
        else ()
    )
    snake_case__ =				(FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()


    def        a				(	self					:					Any       )				->	int:
         lowerCAmelCase__          =			FlaxBlenderbotModelTester(self       )


    def        a				(	self					:					List[str]       )				->	Tuple:
         lowerCAmelCase__          =			self.model_tester.prepare_config_and_inputs()
         for model_class in self.all_model_classes:
              self.model_tester.check_use_cache_forward(_A       ,     _A       ,     _A       )


    def        a				(	self					:					Tuple       )				->	List[str]:
         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(_A       ,     _A       ,     _A       )


    def        a				(	self					:					Optional[Any]       )				->	str:
         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(_A       ,     _A       )
                   lowerCAmelCase__          =			model_class(_A       )

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

                   with self.subTest("JIT Enabled"       ):
                        lowerCAmelCase__          =			encode_jitted(**_A       ).to_tuple()

                   with self.subTest("JIT Disabled"       ):
                        with jax.disable_jit():
                             lowerCAmelCase__          =			encode_jitted(**_A       ).to_tuple()

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


    def        a				(	self					:					Optional[int]       )				->	Union[str, 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(_A       )
                   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(SCREAMING_SNAKE_CASE__					:					Optional[int]       ,     SCREAMING_SNAKE_CASE__					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					int       ):
                        return model.decode(
                            decoder_input_ids=_A       ,     decoder_attention_mask=_A       ,     encoder_outputs=_A       ,     )

                   with self.subTest("JIT Enabled"       ):
                        lowerCAmelCase__          =			decode_jitted(**_A       ).to_tuple()

                   with self.subTest("JIT Disabled"       ):
                        with jax.disable_jit():
                             lowerCAmelCase__          =			decode_jitted(**_A       ).to_tuple()

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


    @slow
    def        a				(	self					:					str       )				->	Union[str, Any]:


         for model_class_name in self.all_model_classes:
              lowerCAmelCase__          =			model_class_name.from_pretrained("facebook/blenderbot-400M-distill"       )
              # FlaxBlenderbotForSequenceClassification expects eos token in input_ids
              lowerCAmelCase__          =			np.ones((1, 1)       ) * model.config.eos_token_id
              lowerCAmelCase__          =			model(_A       )
              self.assertIsNotNone(_A       )


    @unittest.skipUnless(jax_device != "cpu"       ,     "3B test too slow on CPU."       )
    @slow
    def        a				(	self					:					List[str]       )				->	int:
         lowerCAmelCase__          =			{'num_beams': 1, 'early_stopping': True, 'min_length': 15, 'max_length': 25}
         lowerCAmelCase__          =			{'skip_special_tokens': True, 'clean_up_tokenization_spaces': True}

         lowerCAmelCase__          =			FlaxBlenderbotForConditionalGeneration.from_pretrained("facebook/blenderbot-3B"       ,     from_pt=_A       )
         lowerCAmelCase__          =			BlenderbotTokenizer.from_pretrained("facebook/blenderbot-3B"       )

         lowerCAmelCase__          =			['Sam']
         lowerCAmelCase__          =			tokenizer(_A       ,     return_tensors="jax"       )

         lowerCAmelCase__          =			model.generate(**_A       ,     **_A       )
         lowerCAmelCase__          =			'Sam is a great name. It means "sun" in Gaelic.'

         lowerCAmelCase__          =			tokenizer.batch_decode(_A       ,     **_A       )
         assert generated_txt[0].strip() == tgt_text






import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu





class 			__lowerCamelCase	(   nn.Module   ):




    """simple docstring"""

    def __init__(	self					:					Dict       )				->	Optional[int]:
         super().__init__()
         lowerCAmelCase__          =			nn.Linear(3       ,     4       )
         lowerCAmelCase__          =			nn.BatchNormad(4       )
         lowerCAmelCase__          =			nn.Linear(4       ,     5       )


    def        a				(	self					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					int       )				->	List[Any]:
         return self.lineara(self.batchnorm(self.lineara(SCREAMING_SNAKE_CASE__       )       )       )






class 			__lowerCamelCase	(   UpperCamelCase__   ):




    """simple docstring"""

    def        a				(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					str       ,     *SCREAMING_SNAKE_CASE__					:					Tuple       ,     **SCREAMING_SNAKE_CASE__					:					Any       )				->	Union[str, Any]:
         return (args[0] + 1,) + args[1:], kwargs






class 			__lowerCamelCase	(   UpperCamelCase__   ):




    """simple docstring"""

    def        a				(	self					:					List[Any]       ,     SCREAMING_SNAKE_CASE__					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					str       )				->	Dict:
         return output + 1






class 			__lowerCamelCase	(   unittest.TestCase   ):




    """simple docstring"""

    def        a				(	self					:					List[str]       )				->	Tuple:
         lowerCAmelCase__          =			ModelForTest()
         lowerCAmelCase__          =			ModelHook()

         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         self.assertEqual(test_model._hf_hook       ,     SCREAMING_SNAKE_CASE__       )
         self.assertTrue(hasattr(SCREAMING_SNAKE_CASE__       ,     "_old_forward"       )       )

         # Check adding the hook did not change the name or the signature
         self.assertEqual(test_model.forward.__name__       ,     "forward"       )
         self.assertListEqual(list(inspect.signature(test_model.forward       ).parameters       )       ,     ["x"]       )

         remove_hook_from_module(SCREAMING_SNAKE_CASE__       )
         self.assertFalse(hasattr(SCREAMING_SNAKE_CASE__       ,     "_hf_hook"       )       )
         self.assertFalse(hasattr(SCREAMING_SNAKE_CASE__       ,     "_old_forward"       )       )


    def        a				(	self					:					Union[str, Any]       )				->	int:
         lowerCAmelCase__          =			ModelForTest()
         lowerCAmelCase__          =			ModelHook()

         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     append=SCREAMING_SNAKE_CASE__       )

         self.assertEqual(isinstance(test_model._hf_hook       ,     SCREAMING_SNAKE_CASE__       )       ,     SCREAMING_SNAKE_CASE__       )
         self.assertEqual(len(test_model._hf_hook.hooks       )       ,     2       )
         self.assertTrue(hasattr(SCREAMING_SNAKE_CASE__       ,     "_old_forward"       )       )

         # Check adding the hook did not change the name or the signature
         self.assertEqual(test_model.forward.__name__       ,     "forward"       )
         self.assertListEqual(list(inspect.signature(test_model.forward       ).parameters       )       ,     ["x"]       )

         remove_hook_from_module(SCREAMING_SNAKE_CASE__       )
         self.assertFalse(hasattr(SCREAMING_SNAKE_CASE__       ,     "_hf_hook"       )       )
         self.assertFalse(hasattr(SCREAMING_SNAKE_CASE__       ,     "_old_forward"       )       )


    def        a				(	self					:					List[str]       )				->	Any:
         lowerCAmelCase__          =			ModelForTest()
         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			test_model(x + 1       )
         lowerCAmelCase__          =			test_model(x + 2       )

         lowerCAmelCase__          =			PreForwardHook()
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     atol=1e-5       )       )

         # Attaching a hook to a model when it already has one replaces, does not chain
         lowerCAmelCase__          =			PreForwardHook()
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     atol=1e-5       )       )

         # You need to use the sequential hook to chain two or more hooks
         lowerCAmelCase__          =			SequentialHook(PreForwardHook()       ,     PreForwardHook()       )
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )
         assert torch.allclose(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     atol=1e-5       )


    def        a				(	self					:					Any       )				->	Union[str, Any]:
         lowerCAmelCase__          =			ModelForTest()
         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			PostForwardHook()
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE__       ,     output + 1       ,     atol=1e-5       )       )

         # Attaching a hook to a model when it already has one replaces, does not chain
         lowerCAmelCase__          =			PostForwardHook()
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE__       ,     output + 1       ,     atol=1e-5       )       )

         # You need to use the sequential hook to chain two or more hooks
         lowerCAmelCase__          =			SequentialHook(PostForwardHook()       ,     PostForwardHook()       )
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )
         assert torch.allclose(SCREAMING_SNAKE_CASE__       ,     output + 2       ,     atol=1e-5       )


    def        a				(	self					:					Optional[int]       )				->	int:
         lowerCAmelCase__          =			ModelForTest()
         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			PostForwardHook()
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE__       ,     output + 1       )       )
         self.assertTrue(outputa.requires_grad       )

         lowerCAmelCase__          =			True
         lowerCAmelCase__          =			test_model(SCREAMING_SNAKE_CASE__       )
         self.assertFalse(outputa.requires_grad       )


    @require_multi_gpu
    def        a				(	self					:					Optional[Any]       )				->	List[str]:
         lowerCAmelCase__          =			ModelForTest()
         # Everything is on CPU
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )

         # This will move each submodule on different devices
         add_hook_to_module(model.lineara       ,     AlignDevicesHook(execution_device=0       )       )
         add_hook_to_module(model.batchnorm       ,     AlignDevicesHook(execution_device=0       )       )
         add_hook_to_module(model.lineara       ,     AlignDevicesHook(execution_device=1       )       )

         self.assertEqual(model.lineara.weight.device       ,     torch.device(0       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device(0       )       )
         self.assertEqual(model.batchnorm.running_mean.device       ,     torch.device(0       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device(1       )       )

         # We can still make a forward pass. The input does not need to be on any particular device
         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(output.device       ,     torch.device(1       )       )

         # We can add a general hook to put back output on same device as input.
         add_hook_to_module(SCREAMING_SNAKE_CASE__       ,     AlignDevicesHook(io_same_device=SCREAMING_SNAKE_CASE__       )       )
         lowerCAmelCase__          =			torch.randn(2       ,     3       ).to(0       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(output.device       ,     torch.device(0       )       )


    def        a				(	self					:					List[str]       )				->	List[str]:
         lowerCAmelCase__          =			ModelForTest()

         # Everything is on CPU
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )

         # This will move each submodule on different devices
         lowerCAmelCase__          =			{"execution_device": 0 if torch.cuda.is_available() else "cpu", "offload": True}

         add_hook_to_module(model.lineara       ,     AlignDevicesHook(**SCREAMING_SNAKE_CASE__       )       )
         add_hook_to_module(model.batchnorm       ,     AlignDevicesHook(**SCREAMING_SNAKE_CASE__       )       )
         add_hook_to_module(model.lineara       ,     AlignDevicesHook(**SCREAMING_SNAKE_CASE__       )       )

         # Parameters have been offloaded, so on the meta device
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         # Buffers are not included in the offload by default, so are on the execution device
         lowerCAmelCase__          =			torch.device(hook_kwargs["execution_device"]       )
         self.assertEqual(model.batchnorm.running_mean.device       ,     SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(output.device       ,     SCREAMING_SNAKE_CASE__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_module(model.lineara       )
         remove_hook_from_module(model.batchnorm       )
         remove_hook_from_module(model.lineara       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )

         # Now test with buffers included in the offload
         lowerCAmelCase__          =			{
             "execution_device": 0 if torch.cuda.is_available() else "cpu",
             "offload": True,
             "offload_buffers": True,
         }

         add_hook_to_module(model.lineara       ,     AlignDevicesHook(**SCREAMING_SNAKE_CASE__       )       )
         add_hook_to_module(model.batchnorm       ,     AlignDevicesHook(**SCREAMING_SNAKE_CASE__       )       )
         add_hook_to_module(model.lineara       ,     AlignDevicesHook(**SCREAMING_SNAKE_CASE__       )       )

         # Parameters have been offloaded, so on the meta device, buffers included
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.running_mean.device       ,     torch.device("meta"       )       )

         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(output.device       ,     SCREAMING_SNAKE_CASE__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_module(model.lineara       )
         remove_hook_from_module(model.batchnorm       )
         remove_hook_from_module(model.lineara       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )


    def        a				(	self					:					Optional[int]       )				->	Union[str, Any]:
         lowerCAmelCase__          =			ModelForTest()

         # Everything is on CPU
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )

         # This will move each submodule on different devices
         lowerCAmelCase__          =			0 if torch.cuda.is_available() else "cpu"
         attach_align_device_hook(SCREAMING_SNAKE_CASE__       ,     execution_device=SCREAMING_SNAKE_CASE__       ,     offload=SCREAMING_SNAKE_CASE__       )

         # Parameters have been offloaded, so on the meta device
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         # Buffers are not included in the offload by default, so are on the execution device
         lowerCAmelCase__          =			torch.device(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(model.batchnorm.running_mean.device       ,     SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(output.device       ,     SCREAMING_SNAKE_CASE__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_submodules(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )

         # Now test with buffers included in the offload
         attach_align_device_hook(SCREAMING_SNAKE_CASE__       ,     execution_device=SCREAMING_SNAKE_CASE__       ,     offload=SCREAMING_SNAKE_CASE__       ,     offload_buffers=SCREAMING_SNAKE_CASE__       )

         # Parameters have been offloaded, so on the meta device, buffers included
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.running_mean.device       ,     torch.device("meta"       )       )

         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(output.device       ,     SCREAMING_SNAKE_CASE__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_submodules(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )


    def        a				(	self					:					Optional[Any]       )				->	str:
         lowerCAmelCase__          =			ModelForTest()

         # Everything is on CPU
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )

         # This will move each submodule on different devices
         lowerCAmelCase__          =			0 if torch.cuda.is_available() else "cpu"
         attach_align_device_hook(
             SCREAMING_SNAKE_CASE__       ,     execution_device=SCREAMING_SNAKE_CASE__       ,     offload=SCREAMING_SNAKE_CASE__       ,     weights_map=model.state_dict()       )

         # Parameters have been offloaded, so on the meta device
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         # Buffers are not included in the offload by default, so are on the execution device
         lowerCAmelCase__          =			torch.device(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(model.batchnorm.running_mean.device       ,     SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(output.device       ,     SCREAMING_SNAKE_CASE__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_submodules(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )

         # Now test with buffers included in the offload
         attach_align_device_hook(
             SCREAMING_SNAKE_CASE__       ,     execution_device=SCREAMING_SNAKE_CASE__       ,     offload=SCREAMING_SNAKE_CASE__       ,     weights_map=model.state_dict()       ,     offload_buffers=SCREAMING_SNAKE_CASE__       ,     )

         # Parameters have been offloaded, so on the meta device, buffers included
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("meta"       )       )
         self.assertEqual(model.batchnorm.running_mean.device       ,     torch.device("meta"       )       )

         lowerCAmelCase__          =			torch.randn(2       ,     3       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(output.device       ,     SCREAMING_SNAKE_CASE__       )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_submodules(SCREAMING_SNAKE_CASE__       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.batchnorm.weight.device       ,     torch.device("cpu"       )       )
         self.assertEqual(model.lineara.weight.device       ,     torch.device("cpu"       )       )
