



'''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				)   ->		Any:
      super().__init__()
      lowerCAmelCase__			:  Union[str, Any]											= nn.Linear(3  ,4				)
      lowerCAmelCase__			:  List[str]											= nn.BatchNormad(4				)
      lowerCAmelCase__			:  int											= nn.Linear(4  ,5				)






  def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase				)   ->		Union[str, Any]:
      return self.lineara(self.batchnorm(self.lineara(__UpperCAmelCase				)				)				)







class 	lowerCAmelCase_(    SCREAMING_SNAKE_CASE_					):
  '''simple docstring'''




  def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase  ,*__UpperCAmelCase  ,**__UpperCAmelCase				)   ->		Dict:
      return (args[0] + 1,) + args[1:], kwargs







class 	lowerCAmelCase_(    SCREAMING_SNAKE_CASE_					):
  '''simple docstring'''




  def 							UpperCAmelCase_     (   self  ,__UpperCAmelCase  ,__UpperCAmelCase				)   ->		Any:
      return output + 1







class 	lowerCAmelCase_(    unittest.TestCase					):
  '''simple docstring'''




  def 							UpperCAmelCase_     (   self				)   ->		Any:
      lowerCAmelCase__			:  int											= ModelForTest()
      lowerCAmelCase__			:  Dict											= ModelHook()

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




  def 							UpperCAmelCase_     (   self				)   ->		List[Any]:
      lowerCAmelCase__			:  Dict											= ModelForTest()
      lowerCAmelCase__			:  Tuple											= ModelHook()

      add_hook_to_module(__UpperCAmelCase  ,__UpperCAmelCase				)
      add_hook_to_module(__UpperCAmelCase  ,__UpperCAmelCase  ,append=__UpperCAmelCase				)

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




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

      lowerCAmelCase__			:  Dict											= PreForwardHook()
      add_hook_to_module(__UpperCAmelCase  ,__UpperCAmelCase				)
      lowerCAmelCase__			:  List[str]											= test_model(__UpperCAmelCase				)
      self.assertTrue(torch.allclose(__UpperCAmelCase  ,__UpperCAmelCase  ,atol=1E-5				)				)

      # Attaching a hook to a model when it already has one replaces, does not chain
      lowerCAmelCase__			:  Dict											= PreForwardHook()
      add_hook_to_module(__UpperCAmelCase  ,__UpperCAmelCase				)
      lowerCAmelCase__			:  Union[str, Any]											= test_model(__UpperCAmelCase				)
      self.assertTrue(torch.allclose(__UpperCAmelCase  ,__UpperCAmelCase  ,atol=1E-5				)				)

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

      lowerCAmelCase__			:  List[str]											= test_model(__UpperCAmelCase				)
      assert torch.allclose(__UpperCAmelCase  ,__UpperCAmelCase  ,atol=1E-5				)




  def 							UpperCAmelCase_     (   self				)   ->		int:
      lowerCAmelCase__			:  Optional[int]											= ModelForTest()
      lowerCAmelCase__			:  Optional[Any]											= torch.randn(2  ,3				)
      lowerCAmelCase__			:  Union[str, Any]											= test_model(__UpperCAmelCase				)

      lowerCAmelCase__			:  str											= PostForwardHook()
      add_hook_to_module(__UpperCAmelCase  ,__UpperCAmelCase				)
      lowerCAmelCase__			:  Union[str, Any]											= test_model(__UpperCAmelCase				)
      self.assertTrue(torch.allclose(__UpperCAmelCase  ,output + 1  ,atol=1E-5				)				)

      # Attaching a hook to a model when it already has one replaces, does not chain
      lowerCAmelCase__			:  List[Any]											= PostForwardHook()
      add_hook_to_module(__UpperCAmelCase  ,__UpperCAmelCase				)
      lowerCAmelCase__			:  Dict											= test_model(__UpperCAmelCase				)
      self.assertTrue(torch.allclose(__UpperCAmelCase  ,output + 1  ,atol=1E-5				)				)

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

      lowerCAmelCase__			:  Any											= test_model(__UpperCAmelCase				)
      assert torch.allclose(__UpperCAmelCase  ,output + 2  ,atol=1E-5				)




  def 							UpperCAmelCase_     (   self				)   ->		int:
      lowerCAmelCase__			:  Any											= ModelForTest()
      lowerCAmelCase__			:  List[str]											= torch.randn(2  ,3				)
      lowerCAmelCase__			:  Optional[int]											= test_model(__UpperCAmelCase				)

      lowerCAmelCase__			:  Dict											= PostForwardHook()
      add_hook_to_module(__UpperCAmelCase  ,__UpperCAmelCase				)
      lowerCAmelCase__			:  Union[str, Any]											= test_model(__UpperCAmelCase				)
      self.assertTrue(torch.allclose(__UpperCAmelCase  ,output + 1				)				)
      self.assertTrue(outputa.requires_grad				)

      lowerCAmelCase__			:  Any											= True
      lowerCAmelCase__			:  Union[str, Any]											= test_model(__UpperCAmelCase				)
      self.assertFalse(outputa.requires_grad				)




  @require_multi_gpu
  def 							UpperCAmelCase_     (   self				)   ->		Optional[Any]:
      lowerCAmelCase__			:  str											= 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__			:  Any											= torch.randn(2  ,3				)
      lowerCAmelCase__			:  List[str]											= model(__UpperCAmelCase				)
      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(__UpperCAmelCase  ,AlignDevicesHook(io_same_device=__UpperCAmelCase				)				)
      lowerCAmelCase__			:  Union[str, Any]											= torch.randn(2  ,3				).to(0				)
      lowerCAmelCase__			:  List[str]											= model(__UpperCAmelCase				)
      self.assertEqual(output.device  ,torch.device(0				)				)




  def 							UpperCAmelCase_     (   self				)   ->		int:
      lowerCAmelCase__			:  List[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
      lowerCAmelCase__			:  Union[str, Any]											= {"""execution_device""": 0 if torch.cuda.is_available() else """cpu""", """offload""": True}

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

      # 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__			:  Optional[int]											= torch.device(hook_kwargs["""execution_device"""]				)
      self.assertEqual(model.batchnorm.running_mean.device  ,__UpperCAmelCase				)

      lowerCAmelCase__			:  Union[str, Any]											= torch.randn(2  ,3				)
      lowerCAmelCase__			:  Tuple											= model(__UpperCAmelCase				)
      self.assertEqual(output.device  ,__UpperCAmelCase				)

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

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

      # 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__			:  Optional[int]											= torch.randn(2  ,3				)
      lowerCAmelCase__			:  int											= model(__UpperCAmelCase				)
      self.assertEqual(output.device  ,__UpperCAmelCase				)

      # 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 							UpperCAmelCase_     (   self				)   ->		Tuple:
      lowerCAmelCase__			:  Dict											= 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__			:  List[Any]											= 0 if torch.cuda.is_available() else """cpu"""
      attach_align_device_hook(__UpperCAmelCase  ,execution_device=__UpperCAmelCase  ,offload=__UpperCAmelCase				)

      # 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__			:  Dict											= torch.device(__UpperCAmelCase				)
      self.assertEqual(model.batchnorm.running_mean.device  ,__UpperCAmelCase				)

      lowerCAmelCase__			:  List[Any]											= torch.randn(2  ,3				)
      lowerCAmelCase__			:  Optional[int]											= model(__UpperCAmelCase				)
      self.assertEqual(output.device  ,__UpperCAmelCase				)

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

      # 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__			:  Dict											= torch.randn(2  ,3				)
      lowerCAmelCase__			:  List[Any]											= model(__UpperCAmelCase				)
      self.assertEqual(output.device  ,__UpperCAmelCase				)

      # Removing hooks loads back the weights in the model.
      remove_hook_from_submodules(__UpperCAmelCase				)
      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 							UpperCAmelCase_     (   self				)   ->		Tuple:
      lowerCAmelCase__			:  Optional[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
      lowerCAmelCase__			:  str											= 0 if torch.cuda.is_available() else """cpu"""
      attach_align_device_hook(
          __UpperCAmelCase  ,execution_device=__UpperCAmelCase  ,offload=__UpperCAmelCase  ,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__			:  List[str]											= torch.device(__UpperCAmelCase				)
      self.assertEqual(model.batchnorm.running_mean.device  ,__UpperCAmelCase				)

      lowerCAmelCase__			:  Tuple											= torch.randn(2  ,3				)
      lowerCAmelCase__			:  List[str]											= model(__UpperCAmelCase				)
      self.assertEqual(output.device  ,__UpperCAmelCase				)

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

      # 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__			:  int											= torch.randn(2  ,3				)
      lowerCAmelCase__			:  Any											= model(__UpperCAmelCase				)
      self.assertEqual(output.device  ,__UpperCAmelCase				)

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





_lowerCAmelCase										=     '''Input must be a string of 8 numbers plus letter'''
_lowerCAmelCase										=     '''TRWAGMYFPDXBNJZSQVHLCKE'''
def        _SCREAMING_SNAKE_CASE  (	UpperCamelCase      ):



    """simple docstring"""





    if not isinstance(UpperCamelCase							,   UpperCamelCase      ):
        lowerCAmelCase__			:  List[Any]											= f"""Expected string as input, found {type(UpperCamelCase      ).__name__}"""
        raise TypeError(UpperCamelCase      )

    lowerCAmelCase__			:  Union[str, Any]											= spanish_id.replace("""-"""							,   """"""      ).upper()
    if len(UpperCamelCase      ) != 9:
        raise ValueError(UpperCamelCase      )

    try:
        lowerCAmelCase__			:  Optional[int]											= int(spanish_id_clean[0:8]      )
        lowerCAmelCase__			:  int											= spanish_id_clean[8]
    except ValueError as ex:
        raise ValueError(UpperCamelCase      ) from ex

    if letter.isdigit():
        raise ValueError(UpperCamelCase      )

    return letter == LOOKUP_LETTERS[number % 23]


if __name__ == "__main__":
 import doctest

 doctest.testmod()


