





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






       """simple docstring"""





       def __init__(							self   :		Tuple     )    ->     str:

            '''simple docstring'''




            super().__init__()
            _UpperCamelCase				    =       nn.Linear(3       , 4     )
            _UpperCamelCase				    =       nn.BatchNormad(4     )
            _UpperCamelCase				    =       nn.Linear(4       , 5     )







       def snake_case__						(							self   :		Any       , lowerCAmelCase__   :		List[Any]     )    ->     Optional[Any]:

            '''simple docstring'''




            return self.lineara(self.batchnorm(self.lineara(_a     )     )     )

class 				__lowerCAmelCase     (    UpperCAmelCase__ ):






       """simple docstring"""





       def snake_case__						(							self   :		int       , lowerCAmelCase__   :		Union[str, Any]       , *lowerCAmelCase__   :		List[Any]       , **lowerCAmelCase__   :		int     )    ->     Dict:

            '''simple docstring'''




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

class 				__lowerCAmelCase     (    UpperCAmelCase__ ):






       """simple docstring"""





       def snake_case__						(							self   :		Any       , lowerCAmelCase__   :		Tuple       , lowerCAmelCase__   :		int     )    ->     Optional[int]:

            '''simple docstring'''




            return output + 1

class 				__lowerCAmelCase     (    unittest.TestCase ):






       """simple docstring"""





       def snake_case__						(							self   :		int     )    ->     Optional[int]:

            '''simple docstring'''




            _UpperCamelCase				    =       ModelForTest()
            _UpperCamelCase				    =       ModelHook()

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





       def snake_case__						(							self   :		str     )    ->     Dict:

            '''simple docstring'''




            _UpperCamelCase				    =       ModelForTest()
            _UpperCamelCase				    =       ModelHook()

            add_hook_to_module(_a       , _a     )
            add_hook_to_module(_a       , _a       , append=_a     )

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





       def snake_case__						(							self   :		List[Any]     )    ->     Tuple:

            '''simple docstring'''




            _UpperCamelCase				    =       ModelForTest()
            _UpperCamelCase				    =       torch.randn(2       , 3     )
            _UpperCamelCase				    =       test_model(x + 1     )
            _UpperCamelCase				    =       test_model(x + 2     )

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

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

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

            _UpperCamelCase				    =       test_model(_a     )
            assert torch.allclose(_a       , _a       , atol=1e-5     )





       def snake_case__						(							self   :		str     )    ->     int:

            '''simple docstring'''




            _UpperCamelCase				    =       ModelForTest()
            _UpperCamelCase				    =       torch.randn(2       , 3     )
            _UpperCamelCase				    =       test_model(_a     )

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

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

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

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





       def snake_case__						(							self   :		Tuple     )    ->     Optional[Any]:

            '''simple docstring'''




            _UpperCamelCase				    =       ModelForTest()
            _UpperCamelCase				    =       torch.randn(2       , 3     )
            _UpperCamelCase				    =       test_model(_a     )

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

            _UpperCamelCase				    =       True
            _UpperCamelCase				    =       test_model(_a     )
            self.assertFalse(outputa.requires_grad     )





       @require_multi_gpu
       def snake_case__						(							self   :		Tuple     )    ->     Optional[Any]:

            '''simple docstring'''




            _UpperCamelCase				    =       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
            _UpperCamelCase				    =       torch.randn(2       , 3     )
            _UpperCamelCase				    =       model(_a     )
            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(_a       , AlignDevicesHook(io_same_device=_a     )     )
            _UpperCamelCase				    =       torch.randn(2       , 3     ).to(0     )
            _UpperCamelCase				    =       model(_a     )
            self.assertEqual(output.device       , torch.device(0     )     )





       def snake_case__						(							self   :		Dict     )    ->     Optional[Any]:

            '''simple docstring'''




            _UpperCamelCase				    =       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
            _UpperCamelCase				    =       {'''execution_device''': 0 if torch.cuda.is_available() else '''cpu''', '''offload''': True}

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

            # 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
            _UpperCamelCase				    =       torch.device(hook_kwargs['''execution_device''']     )
            self.assertEqual(model.batchnorm.running_mean.device       , _a     )

            _UpperCamelCase				    =       torch.randn(2       , 3     )
            _UpperCamelCase				    =       model(_a     )
            self.assertEqual(output.device       , _a     )

            # 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
            _UpperCamelCase				    =       {
                '''execution_device''': 0 if torch.cuda.is_available() else '''cpu''',
                '''offload''': True,
                '''offload_buffers''': True,
            }

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

            # 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'''     )     )

            _UpperCamelCase				    =       torch.randn(2       , 3     )
            _UpperCamelCase				    =       model(_a     )
            self.assertEqual(output.device       , _a     )

            # 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 snake_case__						(							self   :		str     )    ->     Union[str, Any]:

            '''simple docstring'''




            _UpperCamelCase				    =       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
            _UpperCamelCase				    =       0 if torch.cuda.is_available() else '''cpu'''
            attach_align_device_hook(_a       , execution_device=_a       , offload=_a     )

            # 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
            _UpperCamelCase				    =       torch.device(_a     )
            self.assertEqual(model.batchnorm.running_mean.device       , _a     )

            _UpperCamelCase				    =       torch.randn(2       , 3     )
            _UpperCamelCase				    =       model(_a     )
            self.assertEqual(output.device       , _a     )

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

            # 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'''     )     )

            _UpperCamelCase				    =       torch.randn(2       , 3     )
            _UpperCamelCase				    =       model(_a     )
            self.assertEqual(output.device       , _a     )

            # Removing hooks loads back the weights in the model.
            remove_hook_from_submodules(_a     )
            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 snake_case__						(							self   :		Optional[int]     )    ->     str:

            '''simple docstring'''




            _UpperCamelCase				    =       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
            _UpperCamelCase				    =       0 if torch.cuda.is_available() else '''cpu'''
            attach_align_device_hook(
                _a       , execution_device=_a       , offload=_a       , 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
            _UpperCamelCase				    =       torch.device(_a     )
            self.assertEqual(model.batchnorm.running_mean.device       , _a     )

            _UpperCamelCase				    =       torch.randn(2       , 3     )
            _UpperCamelCase				    =       model(_a     )
            self.assertEqual(output.device       , _a     )

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

            # 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'''     )     )

            _UpperCamelCase				    =       torch.randn(2       , 3     )
            _UpperCamelCase				    =       model(_a     )
            self.assertEqual(output.device       , _a     )

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









"""simple docstring"""




import unittest

import numpy as np

from transformers import RoFormerConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
					import jax.numpy as jnp

					from transformers.models.roformer.modeling_flax_roformer import (
					    FlaxRoFormerForMaskedLM,
					    FlaxRoFormerForMultipleChoice,
					    FlaxRoFormerForQuestionAnswering,
					    FlaxRoFormerForSequenceClassification,
					    FlaxRoFormerForTokenClassification,
					    FlaxRoFormerModel,
					)






class     __magic_name__							(       unittest.TestCase				):






				'''simple docstring'''





				def __init__(       self		, _a		, _a=13		, _a=7		, _a=True		, _a=True		, _a=True		, _a=True		, _a=99		, _a=32		, _a=5		, _a=4		, _a=37		, _a="gelu"		, _a=0.1		, _a=0.1		, _a=512		, _a=16		, _a=2		, _a=0.02		, _a=4		, ):





							"""simple docstring"""
							lowerCamelCase			      =					parent
							lowerCamelCase			      =					batch_size
							lowerCamelCase			      =					seq_length
							lowerCamelCase			      =					is_training
							lowerCamelCase			      =					use_attention_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_choices





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					ids_tensor([self.batch_size, self.seq_length]		, self.vocab_size )

							lowerCamelCase			      =					None
							if self.use_attention_mask:
										lowerCamelCase			      =					random_attention_mask([self.batch_size, self.seq_length] )

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

							lowerCamelCase			      =					RoFormerConfig(
							    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=_a		, initializer_range=self.initializer_range		, )

							return config, input_ids, token_type_ids, attention_mask





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					self.prepare_config_and_inputs()
							lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase			      =					config_and_inputs
							lowerCamelCase			      =					{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
							return config, inputs_dict






@require_flax
class     __magic_name__							(       UpperCAmelCase__						,					unittest.TestCase				):






				'''simple docstring'''
				__UpperCamelCase			  =						True

				__UpperCamelCase			  =						(
				    (
				        FlaxRoFormerModel,
				        FlaxRoFormerForMaskedLM,
				        FlaxRoFormerForSequenceClassification,
				        FlaxRoFormerForTokenClassification,
				        FlaxRoFormerForMultipleChoice,
				        FlaxRoFormerForQuestionAnswering,
				    )
				    if is_flax_available()
				    else ()
				)





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					FlaxRoFormerModelTester(self )





				@slow
				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							for model_class_name in self.all_model_classes:
										lowerCamelCase			      =					model_class_name.from_pretrained("""junnyu/roformer_chinese_small"""		, from_pt=_a )
										lowerCamelCase			      =					model(np.ones((1, 1) ) )
										self.assertIsNotNone(_a )






@require_flax
class     __magic_name__							(       unittest.TestCase				):






				'''simple docstring'''





				@slow
				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					FlaxRoFormerForMaskedLM.from_pretrained("""junnyu/roformer_chinese_base""" )
							lowerCamelCase			      =					jnp.array([[0, 1, 2, 3, 4, 5]] )
							lowerCamelCase			      =					model(_a )[0]

							lowerCamelCase			      =					50_000

							lowerCamelCase			      =					(1, 6, vocab_size)
							self.assertEqual(output.shape		, _a )

							lowerCamelCase			      =					jnp.array(
							    [[[-0.1_205, -1.0_265, 0.2_922], [-1.5_134, 0.1_974, 0.1_519], [-5.0_135, -3.9_003, -0.8_404]]] )

							self.assertTrue(jnp.allclose(output[:, :3, :3]		, _a		, atol=1e-4 ) )



