






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






												"""simple docstring"""
												super().__init__()
												lowerCAmelCase__        =  nn.Linear(3       ,  4				)
												lowerCAmelCase__        =  nn.BatchNormad(4				)
												lowerCAmelCase__        =  nn.Linear(4       ,  5				)







						def  UpperCamelCase__ (		self       ,  _UpperCamelCase				):






												"""simple docstring"""
												return self.lineara(self.batchnorm(self.lineara(_UpperCamelCase				)				)				)




class       __SCREAMING_SNAKE_CASE   (				__lowercase):






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase       ,  *_UpperCamelCase       ,  **_UpperCamelCase				):






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




class       __SCREAMING_SNAKE_CASE   (				__lowercase):






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase       ,  _UpperCamelCase				):






												"""simple docstring"""
												return output + 1




class       __SCREAMING_SNAKE_CASE   (				unittest.TestCase):






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  ModelForTest()
												lowerCAmelCase__        =  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				):






												"""simple docstring"""
												lowerCAmelCase__        =  ModelForTest()
												lowerCAmelCase__        =  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				):






												"""simple docstring"""
												lowerCAmelCase__        =  ModelForTest()
												lowerCAmelCase__        =  torch.randn(2       ,  3				)
												lowerCAmelCase__        =  test_model(x + 1				)
												lowerCAmelCase__        =  test_model(x + 2				)

												lowerCAmelCase__        =  PreForwardHook()
												add_hook_to_module(_UpperCamelCase       ,  _UpperCamelCase				)
												lowerCAmelCase__        =  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__        =  PreForwardHook()
												add_hook_to_module(_UpperCamelCase       ,  _UpperCamelCase				)
												lowerCAmelCase__        =  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__        =  SequentialHook(PreForwardHook()       ,  PreForwardHook()				)
												add_hook_to_module(_UpperCamelCase       ,  _UpperCamelCase				)

												lowerCAmelCase__        =  test_model(_UpperCamelCase				)
												assert torch.allclose(_UpperCamelCase       ,  _UpperCamelCase       ,  atol=1E-5				)






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  ModelForTest()
												lowerCAmelCase__        =  torch.randn(2       ,  3				)
												lowerCAmelCase__        =  test_model(_UpperCamelCase				)

												lowerCAmelCase__        =  PostForwardHook()
												add_hook_to_module(_UpperCamelCase       ,  _UpperCamelCase				)
												lowerCAmelCase__        =  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__        =  PostForwardHook()
												add_hook_to_module(_UpperCamelCase       ,  _UpperCamelCase				)
												lowerCAmelCase__        =  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__        =  SequentialHook(PostForwardHook()       ,  PostForwardHook()				)
												add_hook_to_module(_UpperCamelCase       ,  _UpperCamelCase				)

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






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  ModelForTest()
												lowerCAmelCase__        =  torch.randn(2       ,  3				)
												lowerCAmelCase__        =  test_model(_UpperCamelCase				)

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

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






						@require_multi_gpu
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  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__        =  torch.randn(2       ,  3				)
												lowerCAmelCase__        =  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__        =  torch.randn(2       ,  3				).to(0				)
												lowerCAmelCase__        =  model(_UpperCamelCase				)
												self.assertEqual(output.device       ,  torch.device(0				)				)






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  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__        =  {'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__        =  torch.device(hook_kwargs['execution_device']				)
												self.assertEqual(model.batchnorm.running_mean.device       ,  _UpperCamelCase				)

												lowerCAmelCase__        =  torch.randn(2       ,  3				)
												lowerCAmelCase__        =  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__        =  {
												    '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__        =  torch.randn(2       ,  3				)
												lowerCAmelCase__        =  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				):






												"""simple docstring"""
												lowerCAmelCase__        =  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__        =  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__        =  torch.device(_UpperCamelCase				)
												self.assertEqual(model.batchnorm.running_mean.device       ,  _UpperCamelCase				)

												lowerCAmelCase__        =  torch.randn(2       ,  3				)
												lowerCAmelCase__        =  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__        =  torch.randn(2       ,  3				)
												lowerCAmelCase__        =  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				):






												"""simple docstring"""
												lowerCAmelCase__        =  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__        =  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__        =  torch.device(_UpperCamelCase				)
												self.assertEqual(model.batchnorm.running_mean.device       ,  _UpperCamelCase				)

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













from collections.abc import Iterator, MutableMapping
from dataclasses import dataclass
from typing import Generic, TypeVar

__snake_case      :      Optional[Any]             =       TypeVar("""KEY""")
__snake_case      :      str             =       TypeVar("""VAL""")


@dataclass(frozen=__lowercase						,   slots=__lowercase)
class       __SCREAMING_SNAKE_CASE   (				Generic[KEY, VAL]):
						_SCREAMING_SNAKE_CASE		:    KEY
						_SCREAMING_SNAKE_CASE		:    VAL




class       __SCREAMING_SNAKE_CASE   (				_Item):






						def __init__(		self				):






												"""simple docstring"""
												super().__init__(_UpperCamelCase       ,  _UpperCamelCase				)







						def __bool__(		self				):






												"""simple docstring"""
												return False




__snake_case      :      int             =       _DeletedItem()


class       __SCREAMING_SNAKE_CASE   (				MutableMapping[KEY, VAL]):






						def __init__(		self       ,  _UpperCamelCase = 8       ,  _UpperCamelCase = 0.75				):






												"""simple docstring"""
												lowerCAmelCase__        =  initial_block_size
												lowerCAmelCase__        =  [None] * initial_block_size
												assert 0.0 < capacity_factor < 1.0
												lowerCAmelCase__        =  capacity_factor
												lowerCAmelCase__        =  0






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase				):






												"""simple docstring"""
												return hash(_UpperCamelCase				) % len(self._buckets				)






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase				):






												"""simple docstring"""
												return (ind + 1) % len(self._buckets				)






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase				):






												"""simple docstring"""
												lowerCAmelCase__        =  self._buckets[ind]
												if not stored:
																		lowerCAmelCase__        =  _Item(_UpperCamelCase       ,  _UpperCamelCase				)
																		self._len += 1
																		return True
												elif stored.key == key:
																		lowerCAmelCase__        =  _Item(_UpperCamelCase       ,  _UpperCamelCase				)
																		return True
												else:
																		return False






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  len(self._buckets				) * self._capacity_factor
												return len(self				) >= int(_UpperCamelCase				)






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												if len(self._buckets				) <= self._initial_block_size:
																		return False
												lowerCAmelCase__        =  len(self._buckets				) * self._capacity_factor / 2
												return len(self				) < limit






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase				):






												"""simple docstring"""
												lowerCAmelCase__        =  self._buckets
												lowerCAmelCase__        =  [None] * new_size
												lowerCAmelCase__        =  0
												for item in old_buckets:
																		if item:
																								self._add_item(item.key       ,  item.val				)






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												self._resize(len(self._buckets				) * 2				)






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												self._resize(len(self._buckets				) // 2				)






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase				):






												"""simple docstring"""
												lowerCAmelCase__        =  self._get_bucket_index(_UpperCamelCase				)
												for _ in range(len(self._buckets				)				):
																		yield ind
																		lowerCAmelCase__        =  self._get_next_ind(_UpperCamelCase				)






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase       ,  _UpperCamelCase				):






												"""simple docstring"""
												for ind in self._iterate_buckets(_UpperCamelCase				):
																		if self._try_set(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase				):
																								break






						def __setitem__(		self       ,  _UpperCamelCase       ,  _UpperCamelCase				):






												"""simple docstring"""
												if self._is_full():
																		self._size_up()

												self._add_item(_UpperCamelCase       ,  _UpperCamelCase				)






						def __delitem__(		self       ,  _UpperCamelCase				):






												"""simple docstring"""
												for ind in self._iterate_buckets(_UpperCamelCase				):
																		lowerCAmelCase__        =  self._buckets[ind]
																		if item is None:
																								raise KeyError(_UpperCamelCase				)
																		if item is _deleted:
																								continue
																		if item.key == key:
																								lowerCAmelCase__        =  _deleted
																								self._len -= 1
																								break
												if self._is_sparse():
																		self._size_down()






						def __getitem__(		self       ,  _UpperCamelCase				):






												"""simple docstring"""
												for ind in self._iterate_buckets(_UpperCamelCase				):
																		lowerCAmelCase__        =  self._buckets[ind]
																		if item is None:
																								break
																		if item is _deleted:
																								continue
																		if item.key == key:
																								return item.val
												raise KeyError(_UpperCamelCase				)






						def __len__(		self				):






												"""simple docstring"""
												return self._len






						def __iter__(		self				):






												"""simple docstring"""
												yield from (item.key for item in self._buckets if item)







						def __repr__(		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  ' ,'.join(
												    F"{item.key}: {item.val}" for item in self._buckets if item				)
												return F"HashMap({val_string})"





