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

 def __init__( self	:							List[str]			)   ->						Union[str, Any]:

  """simple docstring"""





  super().__init__()
  UpperCamelCase_					:						Tuple =					nn.Linear(3   ,      4			)
  UpperCamelCase_					:						List[Any] =					nn.BatchNormad(4			)
  UpperCamelCase_					:						str =					nn.Linear(4   ,      5			)


 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[int]   ,      snake_case	:							List[Any]			)   ->						Any:

  """simple docstring"""





  return self.lineara(self.batchnorm(self.lineara(lowercase_			)			)			)





class 		_lowercase						(					_UpperCAmelCase       ):

 def 						SCREAMING_SNAKE_CASE__   ( self	:							List[Any]   ,      snake_case	:							List[str]   ,      *snake_case	:							Optional[int]   ,      **snake_case	:							Tuple			)   ->						Any:

  """simple docstring"""





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





class 		_lowercase						(					_UpperCAmelCase       ):

 def 						SCREAMING_SNAKE_CASE__   ( self	:							List[str]   ,      snake_case	:							int   ,      snake_case	:							Optional[Any]			)   ->						List[str]:

  """simple docstring"""





  return output + 1





class 		_lowercase						(					unittest.TestCase       ):

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Union[str, Any]			)   ->						Tuple:

  """simple docstring"""





  UpperCamelCase_					:						Optional[Any] =					ModelForTest()
  UpperCamelCase_					:						List[str] =					ModelHook()

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

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[Any]			)   ->						str:

  """simple docstring"""





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

  add_hook_to_module(lowercase_   ,      lowercase_			)
  add_hook_to_module(lowercase_   ,      lowercase_   ,      append=lowercase_			)

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

 def 						SCREAMING_SNAKE_CASE__   ( self	:							str			)   ->						Any:

  """simple docstring"""





  UpperCamelCase_					:						str =					ModelForTest()
  UpperCamelCase_					:						Optional[Any] =					torch.randn(2   ,      3			)
  UpperCamelCase_					:						str =					test_model(x + 1			)
  UpperCamelCase_					:						List[str] =					test_model(x + 2			)

  UpperCamelCase_					:						Dict =					PreForwardHook()
  add_hook_to_module(lowercase_   ,      lowercase_			)
  UpperCamelCase_					:						Dict =					test_model(lowercase_			)
  self.assertTrue(torch.allclose(lowercase_   ,      lowercase_   ,      atol=1e-5			)			)

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

  # You need to use the sequential hook to chain two or more hooks
  UpperCamelCase_					:						Dict =					SequentialHook(PreForwardHook()   ,      PreForwardHook()			)
  add_hook_to_module(lowercase_   ,      lowercase_			)

  UpperCamelCase_					:						Any =					test_model(lowercase_			)
  assert torch.allclose(lowercase_   ,      lowercase_   ,      atol=1e-5			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[Any]			)   ->						Optional[Any]:

  """simple docstring"""





  UpperCamelCase_					:						int =					ModelForTest()
  UpperCamelCase_					:						List[Any] =					torch.randn(2   ,      3			)
  UpperCamelCase_					:						Optional[Any] =					test_model(lowercase_			)

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

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

  # You need to use the sequential hook to chain two or more hooks
  UpperCamelCase_					:						Any =					SequentialHook(PostForwardHook()   ,      PostForwardHook()			)
  add_hook_to_module(lowercase_   ,      lowercase_			)

  UpperCamelCase_					:						int =					test_model(lowercase_			)
  assert torch.allclose(lowercase_   ,      output + 2   ,      atol=1e-5			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Dict			)   ->						Union[str, Any]:

  """simple docstring"""





  UpperCamelCase_					:						List[str] =					ModelForTest()
  UpperCamelCase_					:						Union[str, Any] =					torch.randn(2   ,      3			)
  UpperCamelCase_					:						int =					test_model(lowercase_			)

  UpperCamelCase_					:						List[str] =					PostForwardHook()
  add_hook_to_module(lowercase_   ,      lowercase_			)
  UpperCamelCase_					:						Dict =					test_model(lowercase_			)
  self.assertTrue(torch.allclose(lowercase_   ,      output + 1			)			)
  self.assertTrue(outputa.requires_grad			)

  UpperCamelCase_					:						Optional[int] =					True
  UpperCamelCase_					:						Dict =					test_model(lowercase_			)
  self.assertFalse(outputa.requires_grad			)

 @require_multi_gpu
 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[int]			)   ->						Any:

  """simple docstring"""





  UpperCamelCase_					:						List[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
  UpperCamelCase_					:						Optional[Any] =					torch.randn(2   ,      3			)
  UpperCamelCase_					:						Dict =					model(lowercase_			)
  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(lowercase_   ,      AlignDevicesHook(io_same_device=lowercase_			)			)
  UpperCamelCase_					:						List[str] =					torch.randn(2   ,      3			).to(0			)
  UpperCamelCase_					:						Optional[Any] =					model(lowercase_			)
  self.assertEqual(output.device   ,      torch.device(0			)			)

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

  """simple docstring"""





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

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

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

  UpperCamelCase_					:						Tuple =					torch.randn(2   ,      3			)
  UpperCamelCase_					:						Optional[Any] =					model(lowercase_			)
  self.assertEqual(output.device   ,      lowercase_			)

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

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

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

  # 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	:							Tuple			)   ->						Any:

  """simple docstring"""





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

  # 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_					:						List[str] =					torch.device(lowercase_			)
  self.assertEqual(model.batchnorm.running_mean.device   ,      lowercase_			)

  UpperCamelCase_					:						Optional[Any] =					torch.randn(2   ,      3			)
  UpperCamelCase_					:						Tuple =					model(lowercase_			)
  self.assertEqual(output.device   ,      lowercase_			)

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

  # 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_					:						Dict =					torch.randn(2   ,      3			)
  UpperCamelCase_					:						str =					model(lowercase_			)
  self.assertEqual(output.device   ,      lowercase_			)

  # Removing hooks loads back the weights in the model.
  remove_hook_from_submodules(lowercase_			)
  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	:							List[Any]			)   ->						Union[str, Any]:

  """simple docstring"""





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

  UpperCamelCase_					:						Optional[Any] =					torch.randn(2   ,      3			)
  UpperCamelCase_					:						Dict =					model(lowercase_			)
  self.assertEqual(output.device   ,      lowercase_			)

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

  # 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_					:						List[str] =					torch.randn(2   ,      3			)
  UpperCamelCase_					:						List[Any] =					model(lowercase_			)
  self.assertEqual(output.device   ,      lowercase_			)

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






from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxSeqaSeqConfigWithPast
from ...utils import logging


lowercase_       		=  logging.get_logger(__name__)

lowercase_       		=  {
    "t5-small": "https://huggingface.co/t5-small/resolve/main/config.json",
    "t5-base": "https://huggingface.co/t5-base/resolve/main/config.json",
    "t5-large": "https://huggingface.co/t5-large/resolve/main/config.json",
    "t5-3b": "https://huggingface.co/t5-3b/resolve/main/config.json",
    "t5-11b": "https://huggingface.co/t5-11b/resolve/main/config.json",
}


class  A						(							_UpperCAmelCase ):







    """simple docstring"""





    lowerCamelCase			=							't5'
    lowerCamelCase			=							['past_key_values']
    lowerCamelCase			=							{'hidden_size': 'd_model', 'num_attention_heads': 'num_heads', 'num_hidden_layers': 'num_layers'}


    def __init__(      self :       Union[str, Any],lowercase_ :       int=3_2_1_2_8,lowercase_ :       int=5_1_2,lowercase_ :       List[str]=6_4,lowercase_ :       Tuple=2_0_4_8,lowercase_ :       Any=6,lowercase_ :       List[str]=None,lowercase_ :       Union[str, Any]=8,lowercase_ :       int=3_2,lowercase_ :       Dict=1_2_8,lowercase_ :       Optional[int]=0.1,lowercase_ :       List[str]=1E-6,lowercase_ :       Tuple=1.0,lowercase_ :       Any="relu",lowercase_ :       Union[str, Any]=True,lowercase_ :       Optional[Any]=True,lowercase_ :       int=0,lowercase_ :       str=1,**lowercase_ :       str,)->    Any:





        '''simple docstring'''







        A__          =   vocab_size
        A__          =   d_model
        A__          =   d_kv
        A__          =   d_ff
        A__          =   num_layers
        A__          =   (
            num_decoder_layers if num_decoder_layers is not None else self.num_layers
        )  # default = symmetry
        A__          =   num_heads
        A__          =   relative_attention_num_buckets
        A__          =   relative_attention_max_distance
        A__          =   dropout_rate
        A__          =   layer_norm_epsilon
        A__          =   initializer_factor
        A__          =   feed_forward_proj
        A__          =   use_cache

        A__          =   self.feed_forward_proj.split('-'		)
        A__          =   act_info[-1]
        A__          =   act_info[0] == 'gated'

        if len(lowercase_		) > 1 and act_info[0] != "gated" or len(lowercase_		) > 2:
            raise ValueError(
                F'`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer.'
                'Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. '
                '\'gated-gelu\' or \'relu\''		)

        # for backwards compatibility
        if feed_forward_proj == "gated-gelu":
            A__          =   'gelu_new'

        super().__init__(
            pad_token_id=lowercase_,eos_token_id=lowercase_,is_encoder_decoder=lowercase_,**lowercase_,)




class  A						(							_UpperCAmelCase ):







    """simple docstring"""





    @property
    def 						snake_case__     (      self :       Tuple		)->    Mapping[str, Mapping[int, str]]:





        '''simple docstring'''







        A__          =   {
            'input_ids': {0: 'batch', 1: 'encoder_sequence'},
            'attention_mask': {0: 'batch', 1: 'encoder_sequence'},
        }
        if self.use_past:
            A__          =   'past_encoder_sequence + sequence'
            A__          =   {0: 'batch'}
            A__          =   {0: 'batch', 1: 'past_decoder_sequence + sequence'}
        else:
            A__          =   {0: 'batch', 1: 'decoder_sequence'}
            A__          =   {0: 'batch', 1: 'decoder_sequence'}

        if self.use_past:
            self.fill_with_past_key_values_(lowercase_,direction='inputs'		)

        return common_inputs


    @property
    def 						snake_case__     (      self :       Any		)->    int:





        '''simple docstring'''







        return 1_3

