






"""simple docstring"""




import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch

def 			lowerCAmelCase__     (     UpperCamelCase__				):  # picklable for multiprocessing



  '''simple docstring'''
  return x.sum()

def 			lowerCAmelCase__     (     UpperCamelCase__				):  # picklable for multiprocessing



  '''simple docstring'''
  return i + 1






@dataclass
class        UpperCamelCase			:
 UpperCamelCase					:       int
 UpperCamelCase					:       str







class        UpperCamelCase			(	snake_case_      ):


 def _lowercase	(     self		:						Any							)					->       List[str]:
   _a  :       Any             =	{}
   _a  :       List[str]             =	[]
   _a  :       Any             =	1
   _a  :       Tuple             =	[1, 2]
   _a  :       Tuple             =	{"""a""": 1, """b""": 2}
   _a  :       List[str]             =	{"""a""": [1, 2], """b""": [3, 4]}
   _a  :       Optional[int]             =	{"""a""": {"""1""": 1}, """b""": 2}
   _a  :       Any             =	{"""a""": 1, """b""": 2, """c""": 3, """d""": 4}
   _a  :       str             =	{}
   _a  :       List[Any]             =	[]
   _a  :       List[Any]             =	2
   _a  :       Any             =	[2, 3]
   _a  :       Optional[int]             =	{"""a""": 2, """b""": 3}
   _a  :       str             =	{"""a""": [2, 3], """b""": [4, 5]}
   _a  :       List[Any]             =	{"""a""": {"""1""": 2}, """b""": 3}
   _a  :       List[str]             =	{"""a""": 2, """b""": 3, """c""": 4, """d""": 5}
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							)							,						UpperCAmelCase__							)

   _a  :       Optional[Any]             =	2
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							,						num_proc=UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							,						num_proc=UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							,						num_proc=UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							,						num_proc=UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							,						num_proc=UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							,						num_proc=UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							,						num_proc=UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							,						num_proc=UpperCAmelCase__							)							,						UpperCAmelCase__							)

   _a  :       str             =	{"""a""": np.eye(2							), """b""": np.zeros(3							), """c""": np.ones(2							)}
   _a  :       Optional[int]             =	{"""a""": 2, """b""": 0, """c""": 2}
   _a  :       Tuple             =	{
       """a""": np.eye(2							).astype(UpperCAmelCase__							),
       """b""": np.zeros(3							).astype(UpperCAmelCase__							),
       """c""": np.ones(2							).astype(UpperCAmelCase__							),
   }
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							,						map_numpy=UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(
       {k: v.tolist() for k, v in map_nested(UpperCAmelCase__							,						UpperCAmelCase__							,						map_numpy=UpperCAmelCase__							).items()}							,						{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}							,						)
   self.assertEqual(map_nested(UpperCAmelCase__							,						UpperCAmelCase__							,						map_numpy=UpperCAmelCase__							,						num_proc=UpperCAmelCase__							)							,						UpperCAmelCase__							)
   self.assertEqual(
       {k: v.tolist() for k, v in map_nested(UpperCAmelCase__							,						UpperCAmelCase__							,						map_numpy=UpperCAmelCase__							,						num_proc=UpperCAmelCase__							).items()}							,						{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}							,						)
   with self.assertRaises(UpperCAmelCase__							):  # can't pickle a local lambda
     map_nested(lambda UpperCAmelCase__							: x + 1							,						UpperCAmelCase__							,						num_proc=UpperCAmelCase__							)


 def _lowercase	(     self		:						Optional[int]							)					->       int:
   _a  :       Dict             =	{"""a""": 1, """b""": 2}
   _a  :       Optional[Any]             =	{"""a""": 3, """b""": 4}
   _a  :       Union[str, Any]             =	{"""a""": 5, """b""": 6}
   _a  :       Optional[int]             =	sorted([("""a""", (1, 3, 5)), ("""b""", (2, 4, 6))]							)
   self.assertEqual(sorted(zip_dict(UpperCAmelCase__							,						UpperCAmelCase__							,						UpperCAmelCase__							)							)							,						UpperCAmelCase__							)







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






   class        UpperCamelCase			:
    UpperCamelCase					:       Dict			=					'''bar'''







   _a  :       List[str]             =	Foo()
   self.assertEqual(foo.my_attr							,						"""bar"""							)
   with temporary_assignment(UpperCAmelCase__							,						"""my_attr"""							,						"""BAR"""							):
     self.assertEqual(foo.my_attr							,						"""BAR"""							)
   self.assertEqual(foo.my_attr							,						"""bar"""							)







@pytest.mark.parametrize(
    """iterable_length, num_proc, expected_num_proc"""     ,   [
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (1_6, 1_6, 1_6),
        (1_6, 1_7, 1_6),
        (1_7, 1_6, 1_6),
    ]     ,   )
def 			lowerCAmelCase__     (     UpperCamelCase__     ,   UpperCamelCase__     ,   UpperCamelCase__				):



  '''simple docstring'''
  with patch("""datasets.utils.py_utils._single_map_nested"""				) as mock_single_map_nested, patch(
      """datasets.parallel.parallel.Pool"""				) as mock_multiprocessing_pool:
    _a  :       int             =	{F"""{i}""": i for i in range(UpperCamelCase__				)}
    _a  :       Any             =	map_nested(lambda UpperCamelCase__				: x + 1_0     ,   UpperCamelCase__     ,   num_proc=UpperCamelCase__     ,   parallel_min_length=1_6				)
    if expected_num_proc == 1:
      assert mock_single_map_nested.called
      assert not mock_multiprocessing_pool.called
    else:
      assert not mock_single_map_nested.called
      assert mock_multiprocessing_pool.called
      assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc






class        UpperCamelCase			(	snake_case_      ):


 @require_tf
 def _lowercase	(     self		:						List[Any]							)					->       Optional[int]:
   import tensorflow as tf
   from tensorflow.keras import layers

   _a  :       str             =	layers.Dense(2							)

   def gen_random_output():
     _a  :       Dict             =	tf.random.uniform((1, 3)							)
     return model(UpperCAmelCase__							).numpy()

   with temp_seed(42							,						set_tensorflow=UpperCAmelCase__							):
     _a  :       Optional[Any]             =	gen_random_output()
   with temp_seed(42							,						set_tensorflow=UpperCAmelCase__							):
     _a  :       Optional[Any]             =	gen_random_output()
   _a  :       Tuple             =	gen_random_output()

   np.testing.assert_equal(UpperCAmelCase__							,						UpperCAmelCase__							)
   self.assertGreater(np.abs(outa - outa							).sum()							,						0							)


 @require_torch
 def _lowercase	(     self		:						Optional[int]							)					->       str:
   import torch

   def gen_random_output():
     _a  :       Optional[int]             =	torch.nn.Linear(3							,						2							)
     _a  :       Union[str, Any]             =	torch.rand(1							,						3							)
     return model(UpperCAmelCase__							).detach().numpy()

   with temp_seed(42							,						set_pytorch=UpperCAmelCase__							):
     _a  :       List[Any]             =	gen_random_output()
   with temp_seed(42							,						set_pytorch=UpperCAmelCase__							):
     _a  :       str             =	gen_random_output()
   _a  :       List[str]             =	gen_random_output()

   np.testing.assert_equal(UpperCAmelCase__							,						UpperCAmelCase__							)
   self.assertGreater(np.abs(outa - outa							).sum()							,						0							)







 def _lowercase	(     self		:						str							)					->       int:
   def gen_random_output():
     return np.random.rand(1							,						3							)

   with temp_seed(42							):
     _a  :       List[str]             =	gen_random_output()
   with temp_seed(42							):
     _a  :       Union[str, Any]             =	gen_random_output()
   _a  :       Optional[Any]             =	gen_random_output()

   np.testing.assert_equal(UpperCAmelCase__							,						UpperCAmelCase__							)
   self.assertGreater(np.abs(outa - outa							).sum()							,						0							)







@pytest.mark.parametrize("""input_data"""     ,   [{}]				)
def 			lowerCAmelCase__     (     UpperCamelCase__				):



  '''simple docstring'''
  _a  :       Dict             =	NestedDataStructure(UpperCamelCase__				).data
  assert output_data == input_data

@pytest.mark.parametrize(
    """data, expected_output"""     ,   [
        ({}, []),
        ([], []),
        ("""foo""", ["""foo"""]),
        (["""foo""", """bar"""], ["""foo""", """bar"""]),
        ([["""foo""", """bar"""]], ["""foo""", """bar"""]),
        ([[["""foo"""], ["""bar"""]]], ["""foo""", """bar"""]),
        ([[["""foo"""], """bar"""]], ["""foo""", """bar"""]),
        ({"""a""": 1, """b""": 2}, [1, 2]),
        ({"""a""": [1, 2], """b""": [3, 4]}, [1, 2, 3, 4]),
        ({"""a""": [[1, 2]], """b""": [[3, 4]]}, [1, 2, 3, 4]),
        ({"""a""": [[1, 2]], """b""": [3, 4]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [[[3], [4]]]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [[3, 4]]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [3, 4]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [3, [4]]}, [1, 2, 3, 4]),
        ({"""a""": {"""1""": 1}, """b""": 2}, [1, 2]),
        ({"""a""": {"""1""": [1]}, """b""": 2}, [1, 2]),
        ({"""a""": {"""1""": [1]}, """b""": [2]}, [1, 2]),
    ]     ,   )
def 			lowerCAmelCase__     (     UpperCamelCase__     ,   UpperCamelCase__				):



  '''simple docstring'''
  _a  :       List[str]             =	NestedDataStructure(UpperCamelCase__				).flatten()
  assert output == expected_output

def 			lowerCAmelCase__     (     ):



  '''simple docstring'''
  _a  :       Tuple             =	A(x=1     ,   y="""foobar"""				)
  _a  :       Union[str, Any]             =	{"""x""": 1, """y""": """foobar"""}
  assert asdict(UpperCamelCase__				) == expected_output

  _a  :       List[Any]             =	{"""a""": {"""b""": A(x=1_0     ,   y="""foo"""				)}, """c""": [A(x=2_0     ,   y="""bar"""				)]}
  _a  :       List[Any]             =	{"""a""": {"""b""": {"""x""": 1_0, """y""": """foo"""}}, """c""": [{"""x""": 2_0, """y""": """bar"""}]}
  assert asdict(UpperCamelCase__				) == expected_output

  with pytest.raises(UpperCamelCase__				):
    asdict([1, A(x=1_0     ,   y="""foo"""				)]				)

def 			lowerCAmelCase__     (     UpperCamelCase__				):



  '''simple docstring'''
  return text.split()

def 			lowerCAmelCase__     (     UpperCamelCase__				):



  '''simple docstring'''
  yield (time.time(), content)
  time.sleep(2				)
  yield (time.time(), content)



def 			lowerCAmelCase__     (     ):



  '''simple docstring'''
  with Pool(2				) as pool:
    _a  :       str             =	list(iflatmap_unordered(UpperCamelCase__     ,   _split_text     ,   kwargs_iterable=[{"""text""": """hello there"""}] * 1_0				)				)
    assert out.count("""hello"""				) == 1_0
    assert out.count("""there"""				) == 1_0
    assert len(UpperCamelCase__				) == 2_0

  # check multiprocess from pathos (uses dill for pickling)
  with multiprocess.Pool(2				) as pool:
    _a  :       List[str]             =	list(iflatmap_unordered(UpperCamelCase__     ,   _split_text     ,   kwargs_iterable=[{"""text""": """hello there"""}] * 1_0				)				)
    assert out.count("""hello"""				) == 1_0
    assert out.count("""there"""				) == 1_0
    assert len(UpperCamelCase__				) == 2_0

  # check that we get items as fast as possible
  with Pool(2				) as pool:
    _a  :       List[str]             =	[]
    for yield_time, content in iflatmap_unordered(
        UpperCamelCase__     ,   _aseconds_generator_of_aitems_with_timing     ,   kwargs_iterable=[{"""content""": """a"""}, {"""content""": """b"""}]				):
      assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
      out.append(UpperCamelCase__				)
    assert out.count("""a"""				) == 2
    assert out.count("""b"""				) == 2
    assert len(UpperCamelCase__				) == 4









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


 def __init__(     self		:						Union[str, Any]							)					->       int:
   super().__init__()
   _a  :       Optional[Any]             =	nn.Linear(3							,						4							)
   _a  :       Tuple             =	nn.BatchNormad(4							)
   _a  :       Dict             =	nn.Linear(4							,						5							)







 def _lowercase	(     self		:						Optional[int]							,						UpperCAmelCase__		:						List[str]							)					->       int:
   return self.lineara(self.batchnorm(self.lineara(UpperCAmelCase__							)							)							)







class        UpperCamelCase			(	snake_case_      ):


 def _lowercase	(     self		:						Any							,						UpperCAmelCase__		:						Any							,						*UpperCAmelCase__		:						List[str]							,						**UpperCAmelCase__		:						Optional[int]							)					->       Optional[int]:
   return (args[0] + 1,) + args[1:], kwargs







class        UpperCamelCase			(	snake_case_      ):


 def _lowercase	(     self		:						Optional[Any]							,						UpperCAmelCase__		:						Tuple							,						UpperCAmelCase__		:						Optional[Any]							)					->       List[str]:
   return output + 1







class        UpperCamelCase			(	unittest.TestCase      ):


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

   _a  :       Union[str, Any]             =	PreForwardHook()
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   _a  :       Tuple             =	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  :       int             =	PreForwardHook()
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   _a  :       str             =	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  :       int             =	SequentialHook(PreForwardHook()							,						PreForwardHook()							)
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)

   _a  :       Tuple             =	test_model(UpperCAmelCase__							)
   assert torch.allclose(UpperCAmelCase__							,						UpperCAmelCase__							,						atol=1E-5							)


 def _lowercase	(     self		:						Tuple							)					->       int:
   _a  :       Tuple             =	ModelForTest()
   _a  :       Union[str, Any]             =	torch.randn(2							,						3							)
   _a  :       Optional[int]             =	test_model(UpperCAmelCase__							)

   _a  :       int             =	PostForwardHook()
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   _a  :       List[str]             =	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  :       List[Any]             =	PostForwardHook()
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   _a  :       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
   _a  :       Any             =	SequentialHook(PostForwardHook()							,						PostForwardHook()							)
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)

   _a  :       Optional[int]             =	test_model(UpperCAmelCase__							)
   assert torch.allclose(UpperCAmelCase__							,						output + 2							,						atol=1E-5							)


 def _lowercase	(     self		:						Dict							)					->       Optional[Any]:
   _a  :       Any             =	ModelForTest()
   _a  :       List[Any]             =	torch.randn(2							,						3							)
   _a  :       Dict             =	test_model(UpperCAmelCase__							)

   _a  :       Any             =	PostForwardHook()
   add_hook_to_module(UpperCAmelCase__							,						UpperCAmelCase__							)
   _a  :       List[str]             =	test_model(UpperCAmelCase__							)
   self.assertTrue(torch.allclose(UpperCAmelCase__							,						output + 1							)							)
   self.assertTrue(outputa.requires_grad							)

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


 @require_multi_gpu
 def _lowercase	(     self		:						Optional[Any]							)					->       str:
   _a  :       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
   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  :       Optional[int]             =	torch.randn(2							,						3							)
   _a  :       Any             =	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  :       str             =	torch.randn(2							,						3							).to(0							)
   _a  :       Union[str, Any]             =	model(UpperCAmelCase__							)
   self.assertEqual(output.device							,						torch.device(0							)							)


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

   _a  :       int             =	torch.randn(2							,						3							)
   _a  :       str             =	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  :       List[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"""							)							)

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

   _a  :       Union[str, Any]             =	torch.randn(2							,						3							)
   _a  :       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"""							)							)

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

   _a  :       Union[str, Any]             =	torch.randn(2							,						3							)
   _a  :       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__							,						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  :       Any             =	torch.randn(2							,						3							)
   _a  :       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"""							)							)

