


import tempfile

import torch

from diffusers import (
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    UniPCMultistepScheduler,
)

from .test_schedulers import SchedulerCommonTest

class       __magic_name__	(     __lowerCAmelCase):
     A:		List[Any]			=      (DEISMultistepScheduler,)
     A:		Union[str, Any]			=      (("num_inference_steps", 2_5),)


     def      UpperCAmelCase__	(      self							:     str							,	**lowerCamelCase__							:     Any			)  ->      List[str]:





      '''simple docstring'''






      UpperCamelCase__					:    int          =   {
          '''num_train_timesteps''': 1000,
          '''beta_start''': 0.0001,
          '''beta_end''': 0.02,
          '''beta_schedule''': '''linear''',
          '''solver_order''': 2,
      }

      config.update(**lowerCamelCase__			)
      return config


     def      UpperCAmelCase__	(      self							:     Optional[int]							,	lowerCamelCase__							:     List[Any]=0							,	**lowerCamelCase__							:     Optional[int]			)  ->      Union[str, Any]:





      '''simple docstring'''






      UpperCamelCase__					:    Optional[Any]          =   dict(self.forward_default_kwargs			)
      UpperCamelCase__					:    Optional[Any]          =   kwargs.pop('''num_inference_steps'''							,	lowerCamelCase__			)
      UpperCamelCase__					:    Optional[Any]          =   self.dummy_sample
      UpperCamelCase__					:    Dict          =   0.1 * sample
      UpperCamelCase__					:    List[str]          =   [residual + 0.2, residual + 0.15, residual + 0.10]

      for scheduler_class in self.scheduler_classes:
       UpperCamelCase__					:    Optional[int]          =   self.get_scheduler_config(**lowerCamelCase__			)
       UpperCamelCase__					:    Optional[int]          =   scheduler_class(**lowerCamelCase__			)
       scheduler.set_timesteps(lowerCamelCase__			)
       # copy over dummy past residuals
       UpperCamelCase__					:    Optional[Any]          =   dummy_past_residuals[: scheduler.config.solver_order]

       with tempfile.TemporaryDirectory() as tmpdirname:
        scheduler.save_config(lowerCamelCase__			)
        UpperCamelCase__					:    int          =   scheduler_class.from_pretrained(lowerCamelCase__			)
        new_scheduler.set_timesteps(lowerCamelCase__			)
        # copy over dummy past residuals
        UpperCamelCase__					:    Tuple          =   dummy_past_residuals[: new_scheduler.config.solver_order]

       UpperCamelCase__       ,				UpperCamelCase__					:    List[Any]          =   sample, sample
       for t in range(lowerCamelCase__							,	time_step + scheduler.config.solver_order + 1			):
        UpperCamelCase__					:    Optional[int]          =   scheduler.step(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__							,	**lowerCamelCase__			).prev_sample
        UpperCamelCase__					:    Dict          =   new_scheduler.step(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__							,	**lowerCamelCase__			).prev_sample

        assert torch.sum(torch.abs(output - new_output			)			) < 1E-5, "Scheduler outputs are not identical"


     def      UpperCAmelCase__	(      self							:     Dict			)  ->      Dict:





      '''simple docstring'''






      pass


     def      UpperCAmelCase__	(      self							:     Union[str, Any]							,	lowerCamelCase__							:     Dict=0							,	**lowerCamelCase__							:     Any			)  ->      Union[str, Any]:





      '''simple docstring'''






      UpperCamelCase__					:    Union[str, Any]          =   dict(self.forward_default_kwargs			)
      UpperCamelCase__					:    str          =   kwargs.pop('''num_inference_steps'''							,	lowerCamelCase__			)
      UpperCamelCase__					:    Optional[Any]          =   self.dummy_sample
      UpperCamelCase__					:    Dict          =   0.1 * sample
      UpperCamelCase__					:    Any          =   [residual + 0.2, residual + 0.15, residual + 0.10]

      for scheduler_class in self.scheduler_classes:
       UpperCamelCase__					:    Dict          =   self.get_scheduler_config()
       UpperCamelCase__					:    Optional[int]          =   scheduler_class(**lowerCamelCase__			)
       scheduler.set_timesteps(lowerCamelCase__			)

       # copy over dummy past residuals (must be after setting timesteps)
       UpperCamelCase__					:    Any          =   dummy_past_residuals[: scheduler.config.solver_order]

       with tempfile.TemporaryDirectory() as tmpdirname:
        scheduler.save_config(lowerCamelCase__			)
        UpperCamelCase__					:    str          =   scheduler_class.from_pretrained(lowerCamelCase__			)
        # copy over dummy past residuals
        new_scheduler.set_timesteps(lowerCamelCase__			)

        # copy over dummy past residual (must be after setting timesteps)
        UpperCamelCase__					:    Union[str, Any]          =   dummy_past_residuals[: new_scheduler.config.solver_order]

       UpperCamelCase__					:    List[Any]          =   scheduler.step(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__							,	**lowerCamelCase__			).prev_sample
       UpperCamelCase__					:    int          =   new_scheduler.step(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__							,	**lowerCamelCase__			).prev_sample

       assert torch.sum(torch.abs(output - new_output			)			) < 1E-5, "Scheduler outputs are not identical"


     def      UpperCAmelCase__	(      self							:     Optional[int]							,	lowerCamelCase__							:     Tuple=None							,	**lowerCamelCase__							:     int			)  ->      List[Any]:





      '''simple docstring'''






      if scheduler is None:
       UpperCamelCase__					:    str          =   self.scheduler_classes[0]
       UpperCamelCase__					:    Any          =   self.get_scheduler_config(**lowerCamelCase__			)
       UpperCamelCase__					:    int          =   scheduler_class(**lowerCamelCase__			)

      UpperCamelCase__					:    Any          =   self.scheduler_classes[0]
      UpperCamelCase__					:    Union[str, Any]          =   self.get_scheduler_config(**lowerCamelCase__			)
      UpperCamelCase__					:    str          =   scheduler_class(**lowerCamelCase__			)

      UpperCamelCase__					:    List[str]          =   10
      UpperCamelCase__					:    List[Any]          =   self.dummy_model()
      UpperCamelCase__					:    Optional[int]          =   self.dummy_sample_deter
      scheduler.set_timesteps(lowerCamelCase__			)

      for i, t in enumerate(scheduler.timesteps			):
       UpperCamelCase__					:    str          =   model(lowerCamelCase__							,	lowerCamelCase__			)
       UpperCamelCase__					:    Dict          =   scheduler.step(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__			).prev_sample

      return sample


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





      '''simple docstring'''






      UpperCamelCase__					:    List[str]          =   dict(self.forward_default_kwargs			)

      UpperCamelCase__					:    Union[str, Any]          =   kwargs.pop('''num_inference_steps'''							,	lowerCamelCase__			)

      for scheduler_class in self.scheduler_classes:
       UpperCamelCase__					:    Any          =   self.get_scheduler_config()
       UpperCamelCase__					:    Any          =   scheduler_class(**lowerCamelCase__			)

       UpperCamelCase__					:    Union[str, Any]          =   self.dummy_sample
       UpperCamelCase__					:    str          =   0.1 * sample

       if num_inference_steps is not None and hasattr(lowerCamelCase__							,	'''set_timesteps'''			):
        scheduler.set_timesteps(lowerCamelCase__			)
       elif num_inference_steps is not None and not hasattr(lowerCamelCase__							,	'''set_timesteps'''			):
        UpperCamelCase__					:    Any          =   num_inference_steps

       # copy over dummy past residuals (must be done after set_timesteps)
       UpperCamelCase__					:    Union[str, Any]          =   [residual + 0.2, residual + 0.15, residual + 0.10]
       UpperCamelCase__					:    Dict          =   dummy_past_residuals[: scheduler.config.solver_order]

       UpperCamelCase__					:    Any          =   scheduler.timesteps[5]
       UpperCamelCase__					:    str          =   scheduler.timesteps[6]

       UpperCamelCase__					:    Any          =   scheduler.step(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__							,	**lowerCamelCase__			).prev_sample
       UpperCamelCase__					:    Union[str, Any]          =   scheduler.step(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__							,	**lowerCamelCase__			).prev_sample

       self.assertEqual(output_a.shape							,	sample.shape			)
       self.assertEqual(output_a.shape							,	output_a.shape			)


     def      UpperCAmelCase__	(      self							:     List[str]			)  ->      Any:





      '''simple docstring'''






      UpperCamelCase__					:    Optional[Any]          =   DEISMultistepScheduler(**self.get_scheduler_config()			)
      UpperCamelCase__					:    Any          =   self.full_loop(scheduler=lowerCamelCase__			)
      UpperCamelCase__					:    Tuple          =   torch.mean(torch.abs(lowerCamelCase__			)			)

      assert abs(result_mean.item() - 0.2_3916			) < 1E-3

      UpperCamelCase__					:    Dict          =   DPMSolverSinglestepScheduler.from_config(scheduler.config			)
      UpperCamelCase__					:    Optional[int]          =   DPMSolverMultistepScheduler.from_config(scheduler.config			)
      UpperCamelCase__					:    Union[str, Any]          =   UniPCMultistepScheduler.from_config(scheduler.config			)
      UpperCamelCase__					:    Any          =   DEISMultistepScheduler.from_config(scheduler.config			)

      UpperCamelCase__					:    Union[str, Any]          =   self.full_loop(scheduler=lowerCamelCase__			)
      UpperCamelCase__					:    Optional[Any]          =   torch.mean(torch.abs(lowerCamelCase__			)			)

      assert abs(result_mean.item() - 0.2_3916			) < 1E-3


     def      UpperCAmelCase__	(      self							:     Dict			)  ->      List[Any]:





      '''simple docstring'''






      for timesteps in [25, 50, 100, 999, 1000]:
       self.check_over_configs(num_train_timesteps=lowerCamelCase__			)


     def      UpperCAmelCase__	(      self							:     Optional[int]			)  ->      Dict:





      '''simple docstring'''






      self.check_over_configs(thresholding=lowerCamelCase__			)
      for order in [1, 2, 3]:
       for solver_type in ["logrho"]:
        for threshold in [0.5, 1.0, 2.0]:
         for prediction_type in ["epsilon", "sample"]:
          self.check_over_configs(
              thresholding=lowerCamelCase__							,	prediction_type=lowerCamelCase__							,	sample_max_value=lowerCamelCase__							,	algorithm_type='''deis'''							,	solver_order=lowerCamelCase__							,	solver_type=lowerCamelCase__							,	)


     def      UpperCAmelCase__	(      self							:     Optional[Any]			)  ->      Optional[int]:





      '''simple docstring'''






      for prediction_type in ["epsilon", "v_prediction"]:
       self.check_over_configs(prediction_type=lowerCamelCase__			)


     def      UpperCAmelCase__	(      self							:     Tuple			)  ->      Optional[int]:





      '''simple docstring'''






      for algorithm_type in ["deis"]:
       for solver_type in ["logrho"]:
        for order in [1, 2, 3]:
         for prediction_type in ["epsilon", "sample"]:
          self.check_over_configs(
              solver_order=lowerCamelCase__							,	solver_type=lowerCamelCase__							,	prediction_type=lowerCamelCase__							,	algorithm_type=lowerCamelCase__							,	)
          UpperCamelCase__					:    Tuple          =   self.full_loop(
              solver_order=lowerCamelCase__							,	solver_type=lowerCamelCase__							,	prediction_type=lowerCamelCase__							,	algorithm_type=lowerCamelCase__							,	)
          assert not torch.isnan(lowerCamelCase__			).any(), "Samples have nan numbers"


     def      UpperCAmelCase__	(      self							:     str			)  ->      List[str]:





      '''simple docstring'''






      self.check_over_configs(lower_order_final=lowerCamelCase__			)
      self.check_over_configs(lower_order_final=lowerCamelCase__			)


     def      UpperCAmelCase__	(      self							:     Tuple			)  ->      Dict:





      '''simple docstring'''






      for num_inference_steps in [1, 2, 3, 5, 10, 50, 100, 999, 1000]:
       self.check_over_forward(num_inference_steps=lowerCamelCase__							,	time_step=0			)


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





      '''simple docstring'''






      UpperCamelCase__					:    Tuple          =   self.full_loop()
      UpperCamelCase__					:    List[Any]          =   torch.mean(torch.abs(lowerCamelCase__			)			)

      assert abs(result_mean.item() - 0.2_3916			) < 1E-3


     def      UpperCAmelCase__	(      self							:     List[Any]			)  ->      List[Any]:





      '''simple docstring'''






      UpperCamelCase__					:    Dict          =   self.full_loop(prediction_type='''v_prediction'''			)
      UpperCamelCase__					:    List[Any]          =   torch.mean(torch.abs(lowerCamelCase__			)			)

      assert abs(result_mean.item() - 0.091			) < 1E-3


     def      UpperCAmelCase__	(      self							:     Any			)  ->      Tuple:





      '''simple docstring'''






      UpperCamelCase__					:    Dict          =   self.scheduler_classes[0]
      UpperCamelCase__					:    Optional[int]          =   self.get_scheduler_config(thresholding=lowerCamelCase__							,	dynamic_thresholding_ratio=0			)
      UpperCamelCase__					:    Dict          =   scheduler_class(**lowerCamelCase__			)

      UpperCamelCase__					:    str          =   10
      UpperCamelCase__					:    Union[str, Any]          =   self.dummy_model()
      UpperCamelCase__					:    Dict          =   self.dummy_sample_deter.half()
      scheduler.set_timesteps(lowerCamelCase__			)

      for i, t in enumerate(scheduler.timesteps			):
       UpperCamelCase__					:    Dict          =   model(lowerCamelCase__							,	lowerCamelCase__			)
       UpperCamelCase__					:    Optional[int]          =   scheduler.step(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__			).prev_sample

      assert sample.dtype == torch.floataa






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


     def __init__(      self							:     Dict			)  ->      str:





      '''simple docstring'''






      super().__init__()
      UpperCamelCase__					:    int          =   nn.Linear(3							,	4			)
      UpperCamelCase__					:    str          =   nn.BatchNormad(4			)
      UpperCamelCase__					:    List[str]          =   nn.Linear(4							,	5			)


     def      UpperCAmelCase__	(      self							:     int							,	lowerCamelCase__							:     Dict			)  ->      Tuple:





      '''simple docstring'''






      return self.lineara(self.batchnorm(self.lineara(lowerCamelCase__			)			)			)



class       __magic_name__	(     __lowerCAmelCase):


     def      UpperCAmelCase__	(      self							:     Union[str, Any]							,	lowerCamelCase__							:     Any							,	*lowerCamelCase__							:     int							,	**lowerCamelCase__							:     List[str]			)  ->      Any:





      '''simple docstring'''






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



class       __magic_name__	(     __lowerCAmelCase):


     def      UpperCAmelCase__	(      self							:     Any							,	lowerCamelCase__							:     Any							,	lowerCamelCase__							:     Any			)  ->      Any:





      '''simple docstring'''






      return output + 1



class       __magic_name__	(     unittest.TestCase):


     def      UpperCAmelCase__	(      self							:     Tuple			)  ->      Tuple:





      '''simple docstring'''






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

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


     def      UpperCAmelCase__	(      self							:     Dict			)  ->      str:





      '''simple docstring'''






      UpperCamelCase__					:    Dict          =   ModelForTest()
      UpperCamelCase__					:    List[str]          =   ModelHook()

      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__							,	append=lowerCamelCase__			)

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


     def      UpperCAmelCase__	(      self							:     List[str]			)  ->      int:





      '''simple docstring'''






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

      UpperCamelCase__					:    Optional[int]          =   PreForwardHook()
      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
      UpperCamelCase__					:    List[Any]          =   test_model(lowerCamelCase__			)
      self.assertTrue(torch.allclose(lowerCamelCase__							,	lowerCamelCase__							,	atol=1E-5			)			)

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

      # You need to use the sequential hook to chain two or more hooks
      UpperCamelCase__					:    Tuple          =   SequentialHook(PreForwardHook()							,	PreForwardHook()			)
      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)

      UpperCamelCase__					:    Any          =   test_model(lowerCamelCase__			)
      assert torch.allclose(lowerCamelCase__							,	lowerCamelCase__							,	atol=1E-5			)


     def      UpperCAmelCase__	(      self							:     List[Any]			)  ->      Optional[int]:





      '''simple docstring'''






      UpperCamelCase__					:    Dict          =   ModelForTest()
      UpperCamelCase__					:    str          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    Any          =   test_model(lowerCamelCase__			)

      UpperCamelCase__					:    int          =   PostForwardHook()
      add_hook_to_module(lowerCamelCase__							,	lowerCamelCase__			)
      UpperCamelCase__					:    Dict          =   test_model(lowerCamelCase__			)
      self.assertTrue(torch.allclose(lowerCamelCase__							,	output + 1							,	atol=1E-5			)			)

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

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

      UpperCamelCase__					:    Tuple          =   test_model(lowerCamelCase__			)
      assert torch.allclose(lowerCamelCase__							,	output + 2							,	atol=1E-5			)


     def      UpperCAmelCase__	(      self							:     List[Any]			)  ->      Dict:





      '''simple docstring'''






      UpperCamelCase__					:    List[Any]          =   ModelForTest()
      UpperCamelCase__					:    Tuple          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    Tuple          =   test_model(lowerCamelCase__			)

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

      UpperCamelCase__					:    Optional[int]          =   True
      UpperCamelCase__					:    Optional[int]          =   test_model(lowerCamelCase__			)
      self.assertFalse(outputa.requires_grad			)


     @require_multi_gpu
     def      UpperCAmelCase__	(      self							:     Any			)  ->      Dict:





      '''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
      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__					:    Tuple          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    str          =   model(lowerCamelCase__			)
      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(lowerCamelCase__							,	AlignDevicesHook(io_same_device=lowerCamelCase__			)			)
      UpperCamelCase__					:    Tuple          =   torch.randn(2							,	3			).to(0			)
      UpperCamelCase__					:    Optional[int]          =   model(lowerCamelCase__			)
      self.assertEqual(output.device							,	torch.device(0			)			)


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





      '''simple docstring'''






      UpperCamelCase__					:    Dict          =   ModelForTest()

      # Everything is on CPU
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.batchnorm.weight.device							,	torch.device('''cpu'''			)			)
      self.assertEqual(model.lineara.weight.device							,	torch.device('''cpu'''			)			)

      # This will move each submodule on different devices
      UpperCamelCase__					:    int          =   {'''execution_device''': 0 if torch.cuda.is_available() else '''cpu''', '''offload''': True}

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

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

      UpperCamelCase__					:    Dict          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    List[str]          =   model(lowerCamelCase__			)
      self.assertEqual(output.device							,	lowerCamelCase__			)

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

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

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

      # 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							:     str			)  ->      Optional[int]:





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

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

      UpperCamelCase__					:    int          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    Tuple          =   model(lowerCamelCase__			)
      self.assertEqual(output.device							,	lowerCamelCase__			)

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

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

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

      UpperCamelCase__					:    List[str]          =   torch.randn(2							,	3			)
      UpperCamelCase__					:    Any          =   model(lowerCamelCase__			)
      self.assertEqual(output.device							,	lowerCamelCase__			)

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

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

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


