





from __future__ import annotations

import unittest

import numpy as np

from transformers import LayoutLMConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

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


if is_tf_available():
		import tensorflow as tf

		from transformers.models.layoutlm.modeling_tf_layoutlm import (
		    TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFLayoutLMForMaskedLM,
		    TFLayoutLMForQuestionAnswering,
		    TFLayoutLMForSequenceClassification,
		    TFLayoutLMForTokenClassification,
		    TFLayoutLMModel,
		)







class 			lowerCamelCase					:


					def __init__(					self	:  List[str]       ,     __snake_case	:  Dict       ,     __snake_case	:  List[Any]=13       ,     __snake_case	:  List[str]=7       ,     __snake_case	:  Optional[int]=True       ,     __snake_case	:  str=True       ,     __snake_case	:  Optional[int]=True       ,     __snake_case	:  str=True       ,     __snake_case	:  str=99       ,     __snake_case	:  str=32       ,     __snake_case	:  List[str]=2       ,     __snake_case	:  List[str]=4       ,     __snake_case	:  Union[str, Any]=37       ,     __snake_case	:  List[str]="gelu"       ,     __snake_case	:  Dict=0.1       ,     __snake_case	:  Any=0.1       ,     __snake_case	:  Tuple=5_12       ,     __snake_case	:  Optional[int]=16       ,     __snake_case	:  int=2       ,     __snake_case	:  Any=0.02       ,     __snake_case	:  Any=3       ,     __snake_case	:  Dict=4       ,     __snake_case	:  List[str]=None       ,     __snake_case	:  Tuple=10_00       ,     ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		parent
											_snake_case:   int  					=		batch_size
											_snake_case:   Optional[int]  					=		seq_length
											_snake_case:   Any  					=		is_training
											_snake_case:   str  					=		use_input_mask
											_snake_case:   Any  					=		use_token_type_ids
											_snake_case:   Dict  					=		use_labels
											_snake_case:   str  					=		vocab_size
											_snake_case:   int  					=		hidden_size
											_snake_case:   str  					=		num_hidden_layers
											_snake_case:   Optional[Any]  					=		num_attention_heads
											_snake_case:   Tuple  					=		intermediate_size
											_snake_case:   Optional[Any]  					=		hidden_act
											_snake_case:   Dict  					=		hidden_dropout_prob
											_snake_case:   int  					=		attention_probs_dropout_prob
											_snake_case:   Tuple  					=		max_position_embeddings
											_snake_case:   Optional[Any]  					=		type_vocab_size
											_snake_case:   Optional[Any]  					=		type_sequence_label_size
											_snake_case:   Union[str, Any]  					=		initializer_range
											_snake_case:   int  					=		num_labels
											_snake_case:   Any  					=		num_choices
											_snake_case:   List[Any]  					=		scope
											_snake_case:   Dict  					=		range_bbox


					def    SCREAMING_SNAKE_CASE_		(					self	:  int       ):





											'''simple docstring'''





											_snake_case:   Union[str, Any]  					=		ids_tensor([self.batch_size, self.seq_length]       ,     self.vocab_size       )

											# convert bbox to numpy since TF does not support item assignment
											_snake_case:   List[str]  					=		ids_tensor([self.batch_size, self.seq_length, 4]       ,     self.range_bbox       ).numpy()
											# Ensure that bbox is legal
											for i in range(bbox.shape[0]       ):
																	for j in range(bbox.shape[1]       ):
																							if bbox[i, j, 3] < bbox[i, j, 1]:
																													_snake_case:   str  					=		bbox[i, j, 3]
																													_snake_case:   List[Any]  					=		bbox[i, j, 1]
																													_snake_case:   List[Any]  					=		t
																							if bbox[i, j, 2] < bbox[i, j, 0]:
																													_snake_case:   Any  					=		bbox[i, j, 2]
																													_snake_case:   Union[str, Any]  					=		bbox[i, j, 0]
																													_snake_case:   str  					=		t
											_snake_case:   Union[str, Any]  					=		tf.convert_to_tensor(__snake_case       )

											_snake_case:   List[Any]  					=		None
											if self.use_input_mask:
																	_snake_case:   int  					=		random_attention_mask([self.batch_size, self.seq_length]       )

											_snake_case:   List[str]  					=		None
											if self.use_token_type_ids:
																	_snake_case:   str  					=		ids_tensor([self.batch_size, self.seq_length]       ,     self.type_vocab_size       )

											_snake_case:   List[Any]  					=		None
											_snake_case:   Union[str, Any]  					=		None
											_snake_case:   Tuple  					=		None
											if self.use_labels:
																	_snake_case:   List[Any]  					=		ids_tensor([self.batch_size]       ,     self.type_sequence_label_size       )
																	_snake_case:   Tuple  					=		ids_tensor([self.batch_size, self.seq_length]       ,     self.num_labels       )
																	_snake_case:   Tuple  					=		ids_tensor([self.batch_size]       ,     self.num_choices       )

											_snake_case:   List[Any]  					=		LayoutLMConfig(
											    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       ,     initializer_range=self.initializer_range       ,     )

											return config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


					def    SCREAMING_SNAKE_CASE_		(					self	:  str       ,     __snake_case	:  List[Any]       ,     __snake_case	:  List[str]       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  int       ,     __snake_case	:  Dict       ,     __snake_case	:  int       ,     __snake_case	:  Any       ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		TFLayoutLMModel(config=__snake_case       )

											_snake_case:   List[Any]  					=		model(__snake_case       ,     __snake_case       ,     attention_mask=__snake_case       ,     token_type_ids=__snake_case       )
											_snake_case:   Union[str, Any]  					=		model(__snake_case       ,     __snake_case       ,     token_type_ids=__snake_case       )
											_snake_case:   Tuple  					=		model(__snake_case       ,     __snake_case       )
											self.parent.assertEqual(result.last_hidden_state.shape       ,     (self.batch_size, self.seq_length, self.hidden_size)       )
											self.parent.assertEqual(result.pooler_output.shape       ,     (self.batch_size, self.hidden_size)       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  str       ,     __snake_case	:  Any       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  Tuple       ,     __snake_case	:  Dict       ,     __snake_case	:  Dict       ,     __snake_case	:  Dict       ,     __snake_case	:  Tuple       ,     __snake_case	:  List[str]       ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		TFLayoutLMForMaskedLM(config=__snake_case       )

											_snake_case:   Optional[Any]  					=		model(__snake_case       ,     __snake_case       ,     attention_mask=__snake_case       ,     token_type_ids=__snake_case       ,     labels=__snake_case       )
											self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.seq_length, self.vocab_size)       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ,     __snake_case	:  List[str]       ,     __snake_case	:  List[str]       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  List[str]       ,     __snake_case	:  List[Any]       ,     __snake_case	:  Dict       ,     __snake_case	:  Any       ,     __snake_case	:  Optional[int]       ):





											'''simple docstring'''





											_snake_case:   Union[str, Any]  					=		self.num_labels
											_snake_case:   Dict  					=		TFLayoutLMForSequenceClassification(config=__snake_case       )

											_snake_case:   Any  					=		model(__snake_case       ,     __snake_case       ,     attention_mask=__snake_case       ,     token_type_ids=__snake_case       )
											self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.num_labels)       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  Tuple       ,     __snake_case	:  Optional[Any]       ,     __snake_case	:  Optional[int]       ,     __snake_case	:  List[Any]       ,     __snake_case	:  List[str]       ,     __snake_case	:  Optional[int]       ,     __snake_case	:  List[Any]       ):





											'''simple docstring'''





											_snake_case:   Any  					=		self.num_labels
											_snake_case:   Union[str, Any]  					=		TFLayoutLMForTokenClassification(config=__snake_case       )

											_snake_case:   Any  					=		model(__snake_case       ,     __snake_case       ,     attention_mask=__snake_case       ,     token_type_ids=__snake_case       ,     labels=__snake_case       )
											self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.seq_length, self.num_labels)       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[Any]       ,     __snake_case	:  List[Any]       ,     __snake_case	:  Optional[Any]       ,     __snake_case	:  int       ,     __snake_case	:  int       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  Any       ,     __snake_case	:  Optional[int]       ,     __snake_case	:  Optional[Any]       ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		TFLayoutLMForQuestionAnswering(config=__snake_case       )

											_snake_case:   List[str]  					=		model(__snake_case       ,     __snake_case       ,     attention_mask=__snake_case       ,     token_type_ids=__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    SCREAMING_SNAKE_CASE_		(					self	:  List[Any]       ):





											'''simple docstring'''





											_snake_case:   int  					=		self.prepare_config_and_inputs()
											(
											    (
											    _snake_case					
											)					,						(
											    _snake_case					
											)					,						(
											    _snake_case					
											)					,						(
											    _snake_case					
											)					,						(
											    _snake_case					
											)					,						(
											    _snake_case					
											)					,						(
											    _snake_case					
											)					,						(
											    _snake_case					
											)					,											
											):   str  					=		config_and_inputs
											_snake_case:   str  					=		{
											    'input_ids': input_ids,
											    'bbox': bbox,
											    'token_type_ids': token_type_ids,
											    'attention_mask': input_mask,
											}
											return config, inputs_dict







@require_tf
class 			lowerCamelCase					(						__UpperCAmelCase     ,							__UpperCAmelCase     ,							unittest.TestCase  ):
					_SCREAMING_SNAKE_CASE							   =    (
					    (
					        TFLayoutLMModel,
					        TFLayoutLMForMaskedLM,
					        TFLayoutLMForTokenClassification,
					        TFLayoutLMForSequenceClassification,
					        TFLayoutLMForQuestionAnswering,
					    )
					    if is_tf_available()
					    else ()
					)
					_SCREAMING_SNAKE_CASE							   =    (
					    {
					        'feature-extraction': TFLayoutLMModel,
					        'fill-mask': TFLayoutLMForMaskedLM,
					        'text-classification': TFLayoutLMForSequenceClassification,
					        'token-classification': TFLayoutLMForTokenClassification,
					        'zero-shot': TFLayoutLMForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					_SCREAMING_SNAKE_CASE							   =    False
					_SCREAMING_SNAKE_CASE							   =    True
					_SCREAMING_SNAKE_CASE							   =    10


					def    SCREAMING_SNAKE_CASE_		(					self	:  str       ):





											'''simple docstring'''





											_snake_case:   int  					=		TFLayoutLMModelTester(self       )
											_snake_case:   int  					=		ConfigTester(self       ,     config_class=__snake_case       ,     hidden_size=37       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[int]       ):





											'''simple docstring'''





											self.config_tester.run_common_tests()


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											_snake_case:   Any  					=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__snake_case       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  str       ):





											'''simple docstring'''





											_snake_case:   Any  					=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*__snake_case       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  int       ):





											'''simple docstring'''





											_snake_case:   Any  					=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_sequence_classification(*__snake_case       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											_snake_case:   Optional[int]  					=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*__snake_case       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ):





											'''simple docstring'''





											_snake_case:   Tuple  					=		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_question_answering(*__snake_case       )


					@slow
					def    SCREAMING_SNAKE_CASE_		(					self	:  List[Any]       ):





											'''simple docstring'''





											for model_name in TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	_snake_case:   Tuple  					=		TFLayoutLMModel.from_pretrained(__snake_case       )
																	self.assertIsNotNone(__snake_case       )




					@unittest.skip('Onnx compliancy broke with TF 2.10'       )
					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ):





											'''simple docstring'''





											pass





def       lowercase_      (		)   ->Tuple:
						# Here we prepare a batch of 2 sequences to test a LayoutLM forward pass on:
						# fmt: off
						_snake_case:   Tuple  					=		tf.convert_to_tensor([[101,1019,1014,1016,1037,1_2849,4747,1004,1_4246,2278,5439,4524,5002,2930,2193,2930,4341,3208,1005,1055,2171,2848,1_1300,3531,102],[101,4070,4034,7020,1024,3058,1015,1013,2861,1013,6070,1_9274,2772,6205,2_7814,1_6147,1_6147,4343,2047,1_0283,1_0969,1_4389,1012,2338,102]]						)  # noqa: E231
						_snake_case:   List[str]  					=		tf.convert_to_tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],]						)  # noqa: E231
						_snake_case:   Union[str, Any]  					=		tf.convert_to_tensor([[[0,0,0,0],[423,237,440,251],[427,272,441,287],[419,115,437,129],[961,885,992,912],[256,38,330,58],[256,38,330,58],[336,42,353,57],[360,39,401,56],[360,39,401,56],[411,39,471,59],[479,41,528,59],[533,39,630,60],[67,113,134,131],[141,115,209,132],[68,149,133,166],[141,149,187,164],[195,148,287,165],[195,148,287,165],[195,148,287,165],[295,148,349,165],[441,149,492,166],[497,149,546,164],[64,201,125,218],[1000,1000,1000,1000]],[[0,0,0,0],[662,150,754,166],[665,199,742,211],[519,213,554,228],[519,213,554,228],[134,433,187,454],[130,467,204,480],[130,467,204,480],[130,467,204,480],[130,467,204,480],[130,467,204,480],[314,469,376,482],[504,684,582,706],[941,825,973,900],[941,825,973,900],[941,825,973,900],[941,825,973,900],[610,749,652,765],[130,659,168,672],[176,657,237,672],[238,657,312,672],[443,653,628,672],[443,653,628,672],[716,301,825,317],[1000,1000,1000,1000]]]						)  # noqa: E231
						_snake_case:   List[Any]  					=		tf.convert_to_tensor([[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]]						)  # noqa: E231
						# these are sequence labels (i.e. at the token level)
						_snake_case:   Dict  					=		tf.convert_to_tensor([[-100,10,10,10,9,1,-100,7,7,-100,7,7,4,2,5,2,8,8,-100,-100,5,0,3,2,-100],[-100,12,12,12,-100,12,10,-100,-100,-100,-100,10,12,9,-100,-100,-100,10,10,10,9,12,-100,10,-100]]						)  # noqa: E231
						# fmt: on

						return input_ids, attention_mask, bbox, token_type_ids, labels







@require_tf
class 			lowerCamelCase					(						unittest.TestCase  ):


					@slow
					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											_snake_case:   Any  					=		TFLayoutLMModel.from_pretrained('microsoft/layoutlm-base-uncased'       )

											_snake_case					,						_snake_case					,						_snake_case					,						_snake_case					,						_snake_case:   str  					=		prepare_layoutlm_batch_inputs()

											# forward pass
											_snake_case:   Tuple  					=		model(input_ids=__snake_case       ,     bbox=__snake_case       ,     attention_mask=__snake_case       ,     token_type_ids=__snake_case       )

											# test the sequence output on [0, :3, :3]
											_snake_case:   Tuple  					=		tf.convert_to_tensor(
											    [[0.1_785, -0.1_947, -0.0_425], [-0.3_254, -0.2_807, 0.2_553], [-0.5_391, -0.3_322, 0.3_364]]       ,     )

											self.assertTrue(np.allclose(outputs.last_hidden_state[0, :3, :3]       ,     __snake_case       ,     atol=1e-3       )       )

											# test the pooled output on [1, :3]
											_snake_case:   Optional[int]  					=		tf.convert_to_tensor([-0.6_580, -0.0_214, 0.8_552]       )

											self.assertTrue(np.allclose(outputs.pooler_output[1, :3]       ,     __snake_case       ,     atol=1e-3       )       )


					@slow
					def    SCREAMING_SNAKE_CASE_		(					self	:  List[Any]       ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		TFLayoutLMForSequenceClassification.from_pretrained('microsoft/layoutlm-base-uncased'       ,     num_labels=2       )

											_snake_case					,						_snake_case					,						_snake_case					,						_snake_case					,						_snake_case:   Union[str, Any]  					=		prepare_layoutlm_batch_inputs()

											# forward pass
											_snake_case:   Union[str, Any]  					=		model(
											    input_ids=__snake_case       ,     bbox=__snake_case       ,     attention_mask=__snake_case       ,     token_type_ids=__snake_case       ,     labels=tf.convert_to_tensor([1, 1]       )       ,     )

											# test whether we get a loss as a scalar
											_snake_case:   str  					=		outputs.loss
											_snake_case:   Union[str, Any]  					=		(2,)
											self.assertEqual(loss.shape       ,     __snake_case       )

											# test the shape of the logits
											_snake_case:   Optional[int]  					=		outputs.logits
											_snake_case:   Optional[Any]  					=		(2, 2)
											self.assertEqual(logits.shape       ,     __snake_case       )


					@slow
					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ):





											'''simple docstring'''





											_snake_case:   Optional[int]  					=		TFLayoutLMForTokenClassification.from_pretrained('microsoft/layoutlm-base-uncased'       ,     num_labels=13       )

											_snake_case					,						_snake_case					,						_snake_case					,						_snake_case					,						_snake_case:   Dict  					=		prepare_layoutlm_batch_inputs()

											# forward pass
											_snake_case:   Union[str, Any]  					=		model(
											    input_ids=__snake_case       ,     bbox=__snake_case       ,     attention_mask=__snake_case       ,     token_type_ids=__snake_case       ,     labels=__snake_case       )

											# test the shape of the logits
											_snake_case:   Union[str, Any]  					=		outputs.logits
											_snake_case:   List[str]  					=		tf.convert_to_tensor((2, 25, 13)       )
											self.assertEqual(logits.shape       ,     __snake_case       )




					@slow
					def    SCREAMING_SNAKE_CASE_		(					self	:  Union[str, Any]       ):





											'''simple docstring'''





											_snake_case:   Union[str, Any]  					=		TFLayoutLMForQuestionAnswering.from_pretrained('microsoft/layoutlm-base-uncased'       )

											_snake_case					,						_snake_case					,						_snake_case					,						_snake_case					,						_snake_case:   int  					=		prepare_layoutlm_batch_inputs()

											# forward pass
											_snake_case:   List[str]  					=		model(input_ids=__snake_case       ,     bbox=__snake_case       ,     attention_mask=__snake_case       ,     token_type_ids=__snake_case       )

											# test the shape of the logits
											_snake_case:   Union[str, Any]  					=		tf.convert_to_tensor((2, 25)       )
											self.assertEqual(outputs.start_logits.shape       ,     __snake_case       )
											self.assertEqual(outputs.end_logits.shape       ,     __snake_case       )












'''simple docstring'''






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


					def __init__(					self	:  Optional[int]       ):





											'''simple docstring'''





											super().__init__()
											_snake_case:   int  					=		nn.Linear(3       ,     4       )
											_snake_case:   Any  					=		nn.BatchNormad(4       )
											_snake_case:   Union[str, Any]  					=		nn.Linear(4       ,     5       )




					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ,     __snake_case	:  List[Any]       ):





											'''simple docstring'''





											return self.lineara(self.batchnorm(self.lineara(__snake_case       )       )       )







class 			lowerCamelCase					(						__UpperCAmelCase  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ,     __snake_case	:  Optional[Any]       ,     *__snake_case	:  Optional[Any]       ,     **__snake_case	:  List[Any]       ):





											'''simple docstring'''





											return (args[0] + 1,) + args[1:], kwargs







class 			lowerCamelCase					(						__UpperCAmelCase  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  Dict       ,     __snake_case	:  Union[str, Any]       ,     __snake_case	:  Tuple       ):





											'''simple docstring'''





											return output + 1







class 			lowerCamelCase					(						unittest.TestCase  ):


					def    SCREAMING_SNAKE_CASE_		(					self	:  Union[str, Any]       ):





											'''simple docstring'''





											_snake_case:   str  					=		ModelForTest()
											_snake_case:   List[str]  					=		ModelHook()

											add_hook_to_module(__snake_case       ,     __snake_case       )
											self.assertEqual(test_model._hf_hook       ,     __snake_case       )
											self.assertTrue(hasattr(__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(__snake_case       )
											self.assertFalse(hasattr(__snake_case       ,     '_hf_hook'       )       )
											self.assertFalse(hasattr(__snake_case       ,     '_old_forward'       )       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[Any]       ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		ModelForTest()
											_snake_case:   Union[str, Any]  					=		ModelHook()

											add_hook_to_module(__snake_case       ,     __snake_case       )
											add_hook_to_module(__snake_case       ,     __snake_case       ,     append=__snake_case       )

											self.assertEqual(isinstance(test_model._hf_hook       ,     __snake_case       )       ,     __snake_case       )
											self.assertEqual(len(test_model._hf_hook.hooks       )       ,     2       )
											self.assertTrue(hasattr(__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(__snake_case       )
											self.assertFalse(hasattr(__snake_case       ,     '_hf_hook'       )       )
											self.assertFalse(hasattr(__snake_case       ,     '_old_forward'       )       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ):





											'''simple docstring'''





											_snake_case:   Dict  					=		ModelForTest()
											_snake_case:   Tuple  					=		torch.randn(2       ,     3       )
											_snake_case:   List[Any]  					=		test_model(x + 1       )
											_snake_case:   int  					=		test_model(x + 2       )

											_snake_case:   List[str]  					=		PreForwardHook()
											add_hook_to_module(__snake_case       ,     __snake_case       )
											_snake_case:   List[Any]  					=		test_model(__snake_case       )
											self.assertTrue(torch.allclose(__snake_case       ,     __snake_case       ,     atol=1e-5       )       )

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

											# You need to use the sequential hook to chain two or more hooks
											_snake_case:   Any  					=		SequentialHook(PreForwardHook()       ,     PreForwardHook()       )
											add_hook_to_module(__snake_case       ,     __snake_case       )

											_snake_case:   int  					=		test_model(__snake_case       )
											assert torch.allclose(__snake_case       ,     __snake_case       ,     atol=1e-5       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  List[str]       ):





											'''simple docstring'''





											_snake_case:   int  					=		ModelForTest()
											_snake_case:   str  					=		torch.randn(2       ,     3       )
											_snake_case:   int  					=		test_model(__snake_case       )

											_snake_case:   Union[str, Any]  					=		PostForwardHook()
											add_hook_to_module(__snake_case       ,     __snake_case       )
											_snake_case:   int  					=		test_model(__snake_case       )
											self.assertTrue(torch.allclose(__snake_case       ,     output + 1       ,     atol=1e-5       )       )

											# Attaching a hook to a model when it already has one replaces, does not chain
											_snake_case:   Tuple  					=		PostForwardHook()
											add_hook_to_module(__snake_case       ,     __snake_case       )
											_snake_case:   Optional[int]  					=		test_model(__snake_case       )
											self.assertTrue(torch.allclose(__snake_case       ,     output + 1       ,     atol=1e-5       )       )

											# You need to use the sequential hook to chain two or more hooks
											_snake_case:   List[str]  					=		SequentialHook(PostForwardHook()       ,     PostForwardHook()       )
											add_hook_to_module(__snake_case       ,     __snake_case       )

											_snake_case:   List[Any]  					=		test_model(__snake_case       )
											assert torch.allclose(__snake_case       ,     output + 2       ,     atol=1e-5       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Optional[Any]       ):





											'''simple docstring'''





											_snake_case:   Optional[Any]  					=		ModelForTest()
											_snake_case:   Dict  					=		torch.randn(2       ,     3       )
											_snake_case:   Tuple  					=		test_model(__snake_case       )

											_snake_case:   Tuple  					=		PostForwardHook()
											add_hook_to_module(__snake_case       ,     __snake_case       )
											_snake_case:   List[str]  					=		test_model(__snake_case       )
											self.assertTrue(torch.allclose(__snake_case       ,     output + 1       )       )
											self.assertTrue(outputa.requires_grad       )

											_snake_case:   Any  					=		True
											_snake_case:   str  					=		test_model(__snake_case       )
											self.assertFalse(outputa.requires_grad       )


					@require_multi_gpu
					def    SCREAMING_SNAKE_CASE_		(					self	:  Tuple       ):





											'''simple docstring'''





											_snake_case:   Any  					=		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
											_snake_case:   int  					=		torch.randn(2       ,     3       )
											_snake_case:   Tuple  					=		model(__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(__snake_case       ,     AlignDevicesHook(io_same_device=__snake_case       )       )
											_snake_case:   Optional[Any]  					=		torch.randn(2       ,     3       ).to(0       )
											_snake_case:   Any  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     torch.device(0       )       )


					def    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											_snake_case:   Union[str, Any]  					=		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
											_snake_case:   Optional[int]  					=		{'execution_device': 0 if torch.cuda.is_available() else 'cpu', 'offload': True}

											add_hook_to_module(model.lineara       ,     AlignDevicesHook(**__snake_case       )       )
											add_hook_to_module(model.batchnorm       ,     AlignDevicesHook(**__snake_case       )       )
											add_hook_to_module(model.lineara       ,     AlignDevicesHook(**__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
											_snake_case:   Optional[int]  					=		torch.device(hook_kwargs['execution_device']       )
											self.assertEqual(model.batchnorm.running_mean.device       ,     __snake_case       )

											_snake_case:   Tuple  					=		torch.randn(2       ,     3       )
											_snake_case:   Union[str, Any]  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __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
											_snake_case:   Optional[Any]  					=		{
											    'execution_device': 0 if torch.cuda.is_available() else 'cpu',
											    'offload': True,
											    'offload_buffers': True,
											}

											add_hook_to_module(model.lineara       ,     AlignDevicesHook(**__snake_case       )       )
											add_hook_to_module(model.batchnorm       ,     AlignDevicesHook(**__snake_case       )       )
											add_hook_to_module(model.lineara       ,     AlignDevicesHook(**__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'       )       )

											_snake_case:   Optional[Any]  					=		torch.randn(2       ,     3       )
											_snake_case:   Union[str, Any]  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __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    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											_snake_case:   Union[str, Any]  					=		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
											_snake_case:   Optional[Any]  					=		0 if torch.cuda.is_available() else 'cpu'
											attach_align_device_hook(__snake_case       ,     execution_device=__snake_case       ,     offload=__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
											_snake_case:   List[str]  					=		torch.device(__snake_case       )
											self.assertEqual(model.batchnorm.running_mean.device       ,     __snake_case       )

											_snake_case:   Optional[Any]  					=		torch.randn(2       ,     3       )
											_snake_case:   List[str]  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_submodules(__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(__snake_case       ,     execution_device=__snake_case       ,     offload=__snake_case       ,     offload_buffers=__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'       )       )

											_snake_case:   Any  					=		torch.randn(2       ,     3       )
											_snake_case:   List[Any]  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_submodules(__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    SCREAMING_SNAKE_CASE_		(					self	:  Any       ):





											'''simple docstring'''





											_snake_case:   Tuple  					=		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
											_snake_case:   Tuple  					=		0 if torch.cuda.is_available() else 'cpu'
											attach_align_device_hook(
											    __snake_case       ,     execution_device=__snake_case       ,     offload=__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
											_snake_case:   str  					=		torch.device(__snake_case       )
											self.assertEqual(model.batchnorm.running_mean.device       ,     __snake_case       )

											_snake_case:   Dict  					=		torch.randn(2       ,     3       )
											_snake_case:   Dict  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_submodules(__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(
											    __snake_case       ,     execution_device=__snake_case       ,     offload=__snake_case       ,     weights_map=model.state_dict()       ,     offload_buffers=__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'       )       )

											_snake_case:   List[Any]  					=		torch.randn(2       ,     3       )
											_snake_case:   List[str]  					=		model(__snake_case       )
											self.assertEqual(output.device       ,     __snake_case       )

											# Removing hooks loads back the weights in the model.
											remove_hook_from_submodules(__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'       )       )





