


'''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 						__SCREAMING_SNAKE_CASE					(      nn.Module ):






     def __init__(			self		)       ->    Tuple:
          super().__init__()
          _a						=		nn.Linear(3     ,							4		)
          _a						=		nn.BatchNormad(4		)
          _a						=		nn.Linear(4     ,							5		)






     def  a_					(			self     ,							__UpperCamelCase		)       ->    Tuple:
          return self.lineara(self.batchnorm(self.lineara(__UpperCamelCase		)		)		)







class 						__SCREAMING_SNAKE_CASE					(      lowerCamelCase__ ):






     def  a_					(			self     ,							__UpperCamelCase     ,							*__UpperCamelCase     ,							**__UpperCamelCase		)       ->    Any:
          return (args[0] + 1,) + args[1:], kwargs







class 						__SCREAMING_SNAKE_CASE					(      lowerCamelCase__ ):






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    List[str]:
          return output + 1







class 						__SCREAMING_SNAKE_CASE					(      unittest.TestCase ):






     def  a_					(			self		)       ->    List[Any]:
          _a						=		ModelForTest()
          _a						=		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  a_					(			self		)       ->    List[Any]:
          _a						=		ModelForTest()
          _a						=		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  a_					(			self		)       ->    Optional[int]:
          _a						=		ModelForTest()
          _a						=		torch.randn(2     ,							3		)
          _a						=		test_model(x + 1		)
          _a						=		test_model(x + 2		)

          _a						=		PreForwardHook()
          add_hook_to_module(__UpperCamelCase     ,							__UpperCamelCase		)
          _a						=		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
          _a						=		PreForwardHook()
          add_hook_to_module(__UpperCamelCase     ,							__UpperCamelCase		)
          _a						=		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
          _a						=		SequentialHook(PreForwardHook()     ,							PreForwardHook()		)
          add_hook_to_module(__UpperCamelCase     ,							__UpperCamelCase		)

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






     def  a_					(			self		)       ->    Dict:
          _a						=		ModelForTest()
          _a						=		torch.randn(2     ,							3		)
          _a						=		test_model(__UpperCamelCase		)

          _a						=		PostForwardHook()
          add_hook_to_module(__UpperCamelCase     ,							__UpperCamelCase		)
          _a						=		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
          _a						=		PostForwardHook()
          add_hook_to_module(__UpperCamelCase     ,							__UpperCamelCase		)
          _a						=		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
          _a						=		SequentialHook(PostForwardHook()     ,							PostForwardHook()		)
          add_hook_to_module(__UpperCamelCase     ,							__UpperCamelCase		)

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






     def  a_					(			self		)       ->    Optional[Any]:
          _a						=		ModelForTest()
          _a						=		torch.randn(2     ,							3		)
          _a						=		test_model(__UpperCamelCase		)

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

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






     @require_multi_gpu
     def  a_					(			self		)       ->    Tuple:
          _a						=		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
          _a						=		torch.randn(2     ,							3		)
          _a						=		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		)		)
          _a						=		torch.randn(2     ,							3		).to(0		)
          _a						=		model(__UpperCamelCase		)
          self.assertEqual(output.device     ,							torch.device(0		)		)






     def  a_					(			self		)       ->    Union[str, Any]:
          _a						=		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
          _a						=		{"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
          _a						=		torch.device(hook_kwargs["execution_device"]		)
          self.assertEqual(model.batchnorm.running_mean.device     ,							__UpperCamelCase		)

          _a						=		torch.randn(2     ,							3		)
          _a						=		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
          _a						=		{
              "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"		)		)

          _a						=		torch.randn(2     ,							3		)
          _a						=		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  a_					(			self		)       ->    Union[str, Any]:
          _a						=		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
          _a						=		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
          _a						=		torch.device(__UpperCamelCase		)
          self.assertEqual(model.batchnorm.running_mean.device     ,							__UpperCamelCase		)

          _a						=		torch.randn(2     ,							3		)
          _a						=		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"		)		)

          _a						=		torch.randn(2     ,							3		)
          _a						=		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  a_					(			self		)       ->    List[Any]:
          _a						=		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
          _a						=		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
          _a						=		torch.device(__UpperCamelCase		)
          self.assertEqual(model.batchnorm.running_mean.device     ,							__UpperCamelCase		)

          _a						=		torch.randn(2     ,							3		)
          _a						=		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"		)		)

          _a						=		torch.randn(2     ,							3		)
          _a						=		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'''

from typing import TYPE_CHECKING

from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowercase__				     =		{
    "configuration_mctct": ["MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP", "MCTCTConfig"],
    "feature_extraction_mctct": ["MCTCTFeatureExtractor"],
    "processing_mctct": ["MCTCTProcessor"],
}


try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowercase__				     =		[
       "MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST",
       "MCTCTForCTC",
       "MCTCTModel",
       "MCTCTPreTrainedModel",
   ]


if TYPE_CHECKING:
   from .configuration_mctct import MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP, MCTCTConfig
   from .feature_extraction_mctct import MCTCTFeatureExtractor
   from .processing_mctct import MCTCTProcessor

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_mctct import MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST, MCTCTForCTC, MCTCTModel, MCTCTPreTrainedModel

else:
   import sys

   lowercase__				     =		_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

