


import unittest

from parameterized import parameterized

from transformers import AutoTokenizer, GPTNeoXConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         GPTNeoXForCausalLM,
         GPTNeoXForQuestionAnswering,
         GPTNeoXForSequenceClassification,
         GPTNeoXForTokenClassification,
         GPTNeoXModel,
     )





class 			__lowerCamelCase	:




    """simple docstring"""

    def __init__(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					Any=13       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]=7       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]=True       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]=True       ,     SCREAMING_SNAKE_CASE__					:					str=True       ,     SCREAMING_SNAKE_CASE__					:					Any=True       ,     SCREAMING_SNAKE_CASE__					:					str=99       ,     SCREAMING_SNAKE_CASE__					:					Tuple=64       ,     SCREAMING_SNAKE_CASE__					:					Any=5       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]=4       ,     SCREAMING_SNAKE_CASE__					:					Dict=37       ,     SCREAMING_SNAKE_CASE__					:					Dict="gelu"       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]=0.1       ,     SCREAMING_SNAKE_CASE__					:					List[str]=0.1       ,     SCREAMING_SNAKE_CASE__					:					str=512       ,     SCREAMING_SNAKE_CASE__					:					List[Any]=16       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]=2       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]=0.02       ,     SCREAMING_SNAKE_CASE__					:					List[str]=3       ,     SCREAMING_SNAKE_CASE__					:					Tuple=4       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]=None       ,     )				->	int:
         lowerCAmelCase__          =			parent
         lowerCAmelCase__          =			batch_size
         lowerCAmelCase__          =			seq_length
         lowerCAmelCase__          =			is_training
         lowerCAmelCase__          =			use_input_mask
         lowerCAmelCase__          =			use_token_type_ids
         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__          =			type_vocab_size
         lowerCAmelCase__          =			type_sequence_label_size
         lowerCAmelCase__          =			initializer_range
         lowerCAmelCase__          =			num_labels
         lowerCAmelCase__          =			num_choices
         lowerCAmelCase__          =			scope
         lowerCAmelCase__          =			vocab_size - 1


    def        a				(	self					:					List[str]       )				->	Union[str, Any]:
         lowerCAmelCase__          =			ids_tensor([self.batch_size, self.seq_length]       ,     self.vocab_size       )

         lowerCAmelCase__          =			None
         if self.use_input_mask:
              lowerCAmelCase__          =			random_attention_mask([self.batch_size, self.seq_length]       )

         lowerCAmelCase__          =			None
         if self.use_labels:
              lowerCAmelCase__          =			ids_tensor([self.batch_size, self.seq_length]       ,     self.num_labels       )

         lowerCAmelCase__          =			self.get_config()

         return config, input_ids, input_mask, token_labels


    def        a				(	self					:					Optional[int]       )				->	int:
         return GPTNeoXConfig(
             vocab_size=self.vocab_size       ,     hidden_size=self.hidden_size       ,     num_hidden_layers=self.num_hidden_layers       ,     num_attention_heads=self.num_attention_heads       ,     intermediate_size=self.intermediate_size       ,     hidden_act=self.hidden_act       ,     hidden_dropout_prob=self.hidden_dropout_prob       ,     attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,     max_position_embeddings=self.max_position_embeddings       ,     type_vocab_size=self.type_vocab_size       ,     is_decoder=SCREAMING_SNAKE_CASE__       ,     initializer_range=self.initializer_range       ,     pad_token_id=self.pad_token_id       ,     )


    def        a				(	self					:					int       )				->	Optional[int]:
         lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__          =			self.prepare_config_and_inputs()

         lowerCAmelCase__          =			True

         return config, input_ids, input_mask, token_labels


    def        a				(	self					:					Optional[int]       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					str       ,     SCREAMING_SNAKE_CASE__					:					List[str]       )				->	str:
         lowerCAmelCase__          =			GPTNeoXModel(config=SCREAMING_SNAKE_CASE__       )
         model.to(SCREAMING_SNAKE_CASE__       )
         model.eval()
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       ,     attention_mask=SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.parent.assertEqual(result.last_hidden_state.shape       ,     (self.batch_size, self.seq_length, self.hidden_size)       )


    def        a				(	self					:					Optional[int]       ,     SCREAMING_SNAKE_CASE__					:					Dict       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]       )				->	List[str]:
         lowerCAmelCase__          =			True
         lowerCAmelCase__          =			GPTNeoXModel(SCREAMING_SNAKE_CASE__       )
         model.to(SCREAMING_SNAKE_CASE__       )
         model.eval()
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       ,     attention_mask=SCREAMING_SNAKE_CASE__       )
         self.parent.assertEqual(result.last_hidden_state.shape       ,     (self.batch_size, self.seq_length, self.hidden_size)       )


    def        a				(	self					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					str       ,     SCREAMING_SNAKE_CASE__					:					Any       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]       )				->	Optional[int]:
         lowerCAmelCase__          =			GPTNeoXForCausalLM(config=SCREAMING_SNAKE_CASE__       )
         model.to(SCREAMING_SNAKE_CASE__       )
         model.eval()
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       ,     attention_mask=SCREAMING_SNAKE_CASE__       ,     labels=SCREAMING_SNAKE_CASE__       )
         self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.seq_length, self.vocab_size)       )


    def        a				(	self					:					str       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]       ,     SCREAMING_SNAKE_CASE__					:					List[Any]       ,     SCREAMING_SNAKE_CASE__					:					int       )				->	Tuple:
         lowerCAmelCase__          =			self.num_labels
         lowerCAmelCase__          =			GPTNeoXForQuestionAnswering(SCREAMING_SNAKE_CASE__       )
         model.to(SCREAMING_SNAKE_CASE__       )
         model.eval()
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       ,     attention_mask=SCREAMING_SNAKE_CASE__       )
         self.parent.assertEqual(result.start_logits.shape       ,     (self.batch_size, self.seq_length)       )
         self.parent.assertEqual(result.end_logits.shape       ,     (self.batch_size, self.seq_length)       )


    def        a				(	self					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]       )				->	List[str]:
         lowerCAmelCase__          =			self.num_labels
         lowerCAmelCase__          =			GPTNeoXForSequenceClassification(SCREAMING_SNAKE_CASE__       )
         model.to(SCREAMING_SNAKE_CASE__       )
         model.eval()
         lowerCAmelCase__          =			ids_tensor([self.batch_size]       ,     self.type_sequence_label_size       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       ,     attention_mask=SCREAMING_SNAKE_CASE__       ,     labels=SCREAMING_SNAKE_CASE__       )
         self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.num_labels)       )


    def        a				(	self					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					Dict       ,     SCREAMING_SNAKE_CASE__					:					Dict       )				->	str:
         lowerCAmelCase__          =			self.num_labels
         lowerCAmelCase__          =			GPTNeoXForTokenClassification(SCREAMING_SNAKE_CASE__       )
         model.to(SCREAMING_SNAKE_CASE__       )
         model.eval()
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       ,     attention_mask=SCREAMING_SNAKE_CASE__       ,     labels=SCREAMING_SNAKE_CASE__       )
         self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.seq_length, self.num_labels)       )


    def        a				(	self					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					str       ,     SCREAMING_SNAKE_CASE__					:					List[Any]       ,     SCREAMING_SNAKE_CASE__					:					Dict       )				->	str:
         lowerCAmelCase__          =			True
         lowerCAmelCase__          =			GPTNeoXForCausalLM(config=SCREAMING_SNAKE_CASE__       )
         model.to(SCREAMING_SNAKE_CASE__       )
         model.eval()

         # first forward pass
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       ,     attention_mask=SCREAMING_SNAKE_CASE__       ,     use_cache=SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			outputs.past_key_values

         # create hypothetical multiple next token and extent to next_input_ids
         lowerCAmelCase__          =			ids_tensor((self.batch_size, 3)       ,     config.vocab_size       )
         lowerCAmelCase__          =			ids_tensor((self.batch_size, 3)       ,     vocab_size=2       )

         # append to next input_ids and
         lowerCAmelCase__          =			torch.cat([input_ids, next_tokens]       ,     dim=-1       )
         lowerCAmelCase__          =			torch.cat([input_mask, next_mask]       ,     dim=-1       )

         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       ,     attention_mask=SCREAMING_SNAKE_CASE__       ,     output_hidden_states=SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			output_from_no_past["hidden_states"][0]
         lowerCAmelCase__          =			model(
             SCREAMING_SNAKE_CASE__       ,     attention_mask=SCREAMING_SNAKE_CASE__       ,     past_key_values=SCREAMING_SNAKE_CASE__       ,     output_hidden_states=SCREAMING_SNAKE_CASE__       ,     )["hidden_states"][0]

         # select random slice
         lowerCAmelCase__          =			ids_tensor((1,)       ,     output_from_past.shape[-1]       ).item()
         lowerCAmelCase__          =			output_from_no_past[:, -3:, random_slice_idx].detach()
         lowerCAmelCase__          =			output_from_past[:, :, random_slice_idx].detach()

         self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]       )

         # test that outputs are equal for slice
         self.parent.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     atol=1e-3       )       )


    def        a				(	self					:					Union[str, Any]       )				->	Union[str, Any]:
         lowerCAmelCase__          =			self.prepare_config_and_inputs()
         lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__          =			config_and_inputs
         lowerCAmelCase__          =			{"input_ids": input_ids, "attention_mask": input_mask}
         return config, inputs_dict






@require_torch
class 			__lowerCamelCase	(   UpperCamelCase__     ,    UpperCamelCase__     ,    UpperCamelCase__     ,    unittest.TestCase   ):




    """simple docstring"""

    snake_case__ =				(
        (
            GPTNeoXModel,
            GPTNeoXForCausalLM,
            GPTNeoXForQuestionAnswering,
            GPTNeoXForSequenceClassification,
            GPTNeoXForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    snake_case__ =				(GPTNeoXForCausalLM,) if is_torch_available() else ()
    snake_case__ =				(
        {
            "feature-extraction": GPTNeoXModel,
            "question-answering": GPTNeoXForQuestionAnswering,
            "text-classification": GPTNeoXForSequenceClassification,
            "text-generation": GPTNeoXForCausalLM,
            "token-classification": GPTNeoXForTokenClassification,
            "zero-shot": GPTNeoXForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    snake_case__ =				False
    snake_case__ =				False
    snake_case__ =				False
    snake_case__ =				False


    def        a				(	self					:					Union[str, Any]       )				->	Dict:
         lowerCAmelCase__          =			GPTNeoXModelTester(self       )
         lowerCAmelCase__          =			ConfigTester(self       ,     config_class=SCREAMING_SNAKE_CASE__       ,     hidden_size=64       ,     num_attention_heads=8       )


    def        a				(	self					:					List[str]       )				->	Dict:
         self.config_tester.run_common_tests()


    def        a				(	self					:					List[Any]       )				->	Any:
         lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__          =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					Optional[Any]       )				->	int:
         lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__          =			self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_model_as_decoder(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					Union[str, Any]       )				->	List[Any]:
         # This regression test was failing with PyTorch < 1.3
         lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__          =			self.model_tester.prepare_config_and_inputs_for_decoder()

         lowerCAmelCase__          =			None

         self.model_tester.create_and_check_model_as_decoder(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					int       )				->	Dict:
         lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__     ,		lowerCAmelCase__          =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_decoder_model_past_large_inputs(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					str       )				->	Optional[Any]:
         lowerCAmelCase__          =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_causal_lm(*SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					Tuple       )				->	Optional[int]:
         lowerCAmelCase__          =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					Dict       )				->	Union[str, Any]:
         lowerCAmelCase__          =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					Dict       )				->	str:
         lowerCAmelCase__          =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE__       )


    @unittest.skip(reason="Feed forward chunking is not implemented"       )
    def        a				(	self					:					str       )				->	str:
         pass


    @parameterized.expand([("linear",), ("dynamic",)]       )
    def        a				(	self					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]       )				->	Union[str, Any]:
         lowerCAmelCase__     ,		lowerCAmelCase__          =			self.model_tester.prepare_config_and_inputs_for_common()
         lowerCAmelCase__          =			ids_tensor([1, 10]       ,     config.vocab_size       )
         lowerCAmelCase__          =			ids_tensor([1, int(config.max_position_embeddings * 1.5       )]       ,     config.vocab_size       )

         set_seed(42       )  # Fixed seed at init time so the two models get the same random weights
         lowerCAmelCase__          =			GPTNeoXModel(SCREAMING_SNAKE_CASE__       )
         original_model.to(SCREAMING_SNAKE_CASE__       )
         original_model.eval()
         lowerCAmelCase__          =			original_model(SCREAMING_SNAKE_CASE__       ).last_hidden_state
         lowerCAmelCase__          =			original_model(SCREAMING_SNAKE_CASE__       ).last_hidden_state

         set_seed(42       )  # Fixed seed at init time so the two models get the same random weights
         lowerCAmelCase__          =			{"type": scaling_type, "factor": 10.0}
         lowerCAmelCase__          =			GPTNeoXModel(SCREAMING_SNAKE_CASE__       )
         scaled_model.to(SCREAMING_SNAKE_CASE__       )
         scaled_model.eval()
         lowerCAmelCase__          =			scaled_model(SCREAMING_SNAKE_CASE__       ).last_hidden_state
         lowerCAmelCase__          =			scaled_model(SCREAMING_SNAKE_CASE__       ).last_hidden_state

         # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
         # maximum sequence length, so the outputs for the short input should match.
         if scaling_type == "dynamic":
              self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     atol=1e-5       )       )
         else:
              self.assertFalse(torch.allclose(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     atol=1e-5       )       )

         # The output should be different for long inputs
         self.assertFalse(torch.allclose(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     atol=1e-5       )       )






@require_torch
class 			__lowerCamelCase	(   unittest.TestCase   ):




    """simple docstring"""

    @slow
    def        a				(	self					:					str       )				->	List[Any]:
         lowerCAmelCase__          =			AutoTokenizer.from_pretrained("EleutherAI/pythia-410m-deduped"       )
         for checkpointing in [True, False]:
              lowerCAmelCase__          =			GPTNeoXForCausalLM.from_pretrained("EleutherAI/pythia-410m-deduped"       )

              if checkpointing:
                   model.gradient_checkpointing_enable()
              else:
                   model.gradient_checkpointing_disable()
              model.to(SCREAMING_SNAKE_CASE__       )

              lowerCAmelCase__          =			tokenizer("My favorite food is"       ,     return_tensors="pt"       ).to(SCREAMING_SNAKE_CASE__       )
              # The hub repo. is updated on 2023-04-04, resulting in poor outputs.
              # See: https://github.com/huggingface/transformers/pull/24193
              lowerCAmelCase__          =			"My favorite food is a good old-fashioned, old-fashioned, old-fashioned.\n\nI'm not sure"

              lowerCAmelCase__          =			model.generate(**SCREAMING_SNAKE_CASE__       ,     do_sample=SCREAMING_SNAKE_CASE__       ,     max_new_tokens=20       )
              lowerCAmelCase__          =			tokenizer.batch_decode(SCREAMING_SNAKE_CASE__       )[0]

              self.assertEqual(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )






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