




"""simple docstring"""


from dataclasses import dataclass
from typing import Dict, Optional, Union

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .attention_processor import AttentionProcessor, AttnProcessor
from .embeddings import TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin






@dataclass
class  lowercase(					__a			):







 '''simple docstring'''





 lowercase__			 = 42







class  lowercase(					__a	,      __a			):







 '''simple docstring'''





 @register_to_config
 def __init__(    self:		List[Any],				a_:		int = 32,				a_:		int = 64,				a_:		int = 20,				a_:		int = 768,				a_:		Union[str, Any]=77,				a_:		Optional[int]=4,				a_:		float = 0.0,				a_:		str = "silu",				a_:		Optional[str] = None,				a_:		Optional[str] = None,				a_:		Optional[str] = "linear",				a_:		Optional[str] = "prd",				a_:		Optional[int] = None,				a_:		Optional[int] = None,				a_:		Optional[int] = None,				):




   '''simple docstring'''



   super().__init__()
   _snake_case							:		Tuple								=  num_attention_heads
   _snake_case							:		List[str]								=  attention_head_dim
   _snake_case							:		Union[str, Any]								=  num_attention_heads * attention_head_dim
   _snake_case							:		List[Any]								=  additional_embeddings

   _snake_case							:		Optional[int]								=  time_embed_dim or inner_dim
   _snake_case							:		Optional[Any]								=  embedding_proj_dim or embedding_dim
   _snake_case							:		Optional[Any]								=  clip_embed_dim or embedding_dim

   _snake_case							:		List[Any]								=  Timesteps(a_,				a_,				0		)
   _snake_case							:		List[Any]								=  TimestepEmbedding(a_,				a_,				out_dim=a_,				act_fn=a_		)

   _snake_case							:		Optional[Any]								=  nn.Linear(a_,				a_		)

   if embedding_proj_norm_type is None:
     _snake_case							:		Union[str, Any]								=  None
   elif embedding_proj_norm_type == "layer":
     _snake_case							:		List[str]								=  nn.LayerNorm(a_		)
   else:
     raise ValueError(f"unsupported embedding_proj_norm_type: {embedding_proj_norm_type}"		)

   _snake_case							:		Any								=  nn.Linear(a_,				a_		)

   if encoder_hid_proj_type is None:
     _snake_case							:		List[str]								=  None
   elif encoder_hid_proj_type == "linear":
     _snake_case							:		Optional[Any]								=  nn.Linear(a_,				a_		)
   else:
     raise ValueError(f"unsupported encoder_hid_proj_type: {encoder_hid_proj_type}"		)

   _snake_case							:		Any								=  nn.Parameter(torch.zeros(1,				num_embeddings + additional_embeddings,				a_		)		)

   if added_emb_type == "prd":
     _snake_case							:		List[str]								=  nn.Parameter(torch.zeros(1,				1,				a_		)		)
   elif added_emb_type is None:
     _snake_case							:		List[Any]								=  None
   else:
     raise ValueError(
         f"`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `'prd'` or `None`."		)

   _snake_case							:		int								=  nn.ModuleList(
       [
           BasicTransformerBlock(
               a_,				a_,				a_,				dropout=a_,				activation_fn="""gelu""",				attention_bias=a_,				)
           for d in range(a_		)
       ]		)

   if norm_in_type == "layer":
     _snake_case							:		Optional[Any]								=  nn.LayerNorm(a_		)
   elif norm_in_type is None:
     _snake_case							:		str								=  None
   else:
     raise ValueError(f"Unsupported norm_in_type: {norm_in_type}."		)

   _snake_case							:		Optional[int]								=  nn.LayerNorm(a_		)

   _snake_case							:		List[Any]								=  nn.Linear(a_,				a_		)

   _snake_case							:		Optional[int]								=  torch.full(
       [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings],				-10_000.0		)
   causal_attention_mask.triu_(1		)
   _snake_case							:		List[Any]								=  causal_attention_mask[None, ...]
   self.register_buffer("""causal_attention_mask""",				a_,				persistent=a_		)

   _snake_case							:		str								=  nn.Parameter(torch.zeros(1,				a_		)		)
   _snake_case							:		str								=  nn.Parameter(torch.zeros(1,				a_		)		)


 @property
 # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
 def 			UpperCamelCase_       (    self:		str		):




   '''simple docstring'''



   _snake_case							:		str								=  {}

   def fn_recursive_add_processors(a_:		str,				a_:		torch.nn.Module,				a_:		Dict[str, AttentionProcessor]		):
     if hasattr(a_,				"""set_processor"""		):
       _snake_case							:		int								=  module.processor

     for sub_name, child in module.named_children():
       fn_recursive_add_processors(f"{name}.{sub_name}",				a_,				a_		)

     return processors

   for name, module in self.named_children():
     fn_recursive_add_processors(a_,				a_,				a_		)

   return processors


 def 			UpperCamelCase_       (    self:		Optional[Any],				a_:		Union[AttentionProcessor, Dict[str, AttentionProcessor]]		):




   '''simple docstring'''



   _snake_case							:		Optional[int]								=  len(self.attn_processors.keys()		)

   if isinstance(a_,				a_		) and len(a_		) != count:
     raise ValueError(
         f"A dict of processors was passed, but the number of processors {len(a_		)} does not match the"
         f" number of attention layers: {count}. Please make sure to pass {count} processor classes."		)

   def fn_recursive_attn_processor(a_:		str,				a_:		torch.nn.Module,				a_:		Dict		):
     if hasattr(a_,				"""set_processor"""		):
       if not isinstance(a_,				a_		):
         module.set_processor(a_		)
       else:
         module.set_processor(processor.pop(f"{name}.processor"		)		)

     for sub_name, child in module.named_children():
       fn_recursive_attn_processor(f"{name}.{sub_name}",				a_,				a_		)

   for name, module in self.named_children():
     fn_recursive_attn_processor(a_,				a_,				a_		)


 def 			UpperCamelCase_       (    self:		Optional[Any]		):




   '''simple docstring'''



   self.set_attn_processor(AttnProcessor()		)


 def 			UpperCamelCase_       (    self:		Tuple,				a_:		Tuple,				a_:		Union[torch.Tensor, float, int],				a_:		torch.FloatTensor,				a_:		Optional[torch.FloatTensor] = None,				a_:		Optional[torch.BoolTensor] = None,				a_:		bool = True,				):




   '''simple docstring'''



   _snake_case							:		Optional[Any]								=  hidden_states.shape[0]

   _snake_case							:		Optional[Any]								=  timestep
   if not torch.is_tensor(a_		):
     _snake_case							:		int								=  torch.tensor([timesteps],				dtype=torch.long,				device=hidden_states.device		)
   elif torch.is_tensor(a_		) and len(timesteps.shape		) == 0:
     _snake_case							:		Tuple								=  timesteps[None].to(hidden_states.device		)

   # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
   _snake_case							:		Optional[int]								=  timesteps * torch.ones(a_,				dtype=timesteps.dtype,				device=timesteps.device		)

   _snake_case							:		int								=  self.time_proj(a_		)

   # timesteps does not contain any weights and will always return f32 tensors
   # but time_embedding might be fp16, so we need to cast here.
   _snake_case							:		List[str]								=  timesteps_projected.to(dtype=self.dtype		)
   _snake_case							:		Optional[Any]								=  self.time_embedding(a_		)

   if self.embedding_proj_norm is not None:
     _snake_case							:		str								=  self.embedding_proj_norm(a_		)

   _snake_case							:		int								=  self.embedding_proj(a_		)
   if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
     _snake_case							:		Any								=  self.encoder_hidden_states_proj(a_		)
   elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None:
     raise ValueError("""`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set"""		)

   _snake_case							:		Dict								=  self.proj_in(a_		)

   _snake_case							:		Optional[int]								=  self.positional_embedding.to(hidden_states.dtype		)

   _snake_case							:		Union[str, Any]								=  []
   _snake_case							:		str								=  0

   if encoder_hidden_states is not None:
     additional_embeds.append(a_		)
     additional_embeddings_len += encoder_hidden_states.shape[1]

   if len(proj_embeddings.shape		) == 2:
     _snake_case							:		Union[str, Any]								=  proj_embeddings[:, None, :]

   if len(hidden_states.shape		) == 2:
     _snake_case							:		str								=  hidden_states[:, None, :]

   _snake_case							:		Union[str, Any]								=  additional_embeds + [
       proj_embeddings,
       time_embeddings[:, None, :],
       hidden_states,
   ]

   if self.prd_embedding is not None:
     _snake_case							:		Union[str, Any]								=  self.prd_embedding.to(hidden_states.dtype		).expand(a_,				-1,				-1		)
     additional_embeds.append(a_		)

   _snake_case							:		Optional[int]								=  torch.cat(
       a_,				dim=1,				)

   # Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
   _snake_case							:		Union[str, Any]								=  additional_embeddings_len + proj_embeddings.shape[1] + 1
   if positional_embeddings.shape[1] < hidden_states.shape[1]:
     _snake_case							:		int								=  F.pad(
         a_,				(
             0,
             0,
             additional_embeddings_len,
             self.prd_embedding.shape[1] if self.prd_embedding is not None else 0,
         ),				value=0.0,				)

   _snake_case							:		Optional[Any]								=  hidden_states + positional_embeddings

   if attention_mask is not None:
     _snake_case							:		str								=  (1 - attention_mask.to(hidden_states.dtype		)) * -10_000.0
     _snake_case							:		Optional[Any]								=  F.pad(a_,				(0, self.additional_embeddings),				value=0.0		)
     _snake_case							:		int								=  (attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype		)
     _snake_case							:		str								=  attention_mask.repeat_interleave(self.config.num_attention_heads,				dim=0		)

   if self.norm_in is not None:
     _snake_case							:		Dict								=  self.norm_in(a_		)

   for block in self.transformer_blocks:
     _snake_case							:		List[Any]								=  block(a_,				attention_mask=a_		)

   _snake_case							:		Tuple								=  self.norm_out(a_		)

   if self.prd_embedding is not None:
     _snake_case							:		Dict								=  hidden_states[:, -1]
   else:
     _snake_case							:		Tuple								=  hidden_states[:, additional_embeddings_len:]

   _snake_case							:		Any								=  self.proj_to_clip_embeddings(a_		)

   if not return_dict:
     return (predicted_image_embedding,)

   return PriorTransformerOutput(predicted_image_embedding=a_		)


 def 			UpperCamelCase_       (    self:		Optional[Any],				a_:		List[str]		):




   '''simple docstring'''



   _snake_case							:		Optional[Any]								=  (prior_latents * self.clip_std) + self.clip_mean
   return prior_latents








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







 '''simple docstring'''





 def __init__(    self:		Optional[int]		):




   '''simple docstring'''



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


 def 			UpperCamelCase_       (    self:		Any,				a_:		Union[str, Any]		):




   '''simple docstring'''



   return self.lineara(self.batchnorm(self.lineara(a_		)		)		)







class  lowercase(					__a			):







 '''simple docstring'''





 def 			UpperCamelCase_       (    self:		Any,				a_:		int,				*a_:		Dict,				**a_:		Dict		):




   '''simple docstring'''



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







class  lowercase(					__a			):







 '''simple docstring'''





 def 			UpperCamelCase_       (    self:		str,				a_:		Tuple,				a_:		Union[str, Any]		):




   '''simple docstring'''



   return output + 1







class  lowercase(					unittest.TestCase			):







 '''simple docstring'''





 def 			UpperCamelCase_       (    self:		Union[str, Any]		):




   '''simple docstring'''



   _snake_case							:		Tuple								=  ModelForTest()
   _snake_case							:		List[str]								=  ModelHook()

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


 def 			UpperCamelCase_       (    self:		Tuple		):




   '''simple docstring'''



   _snake_case							:		Dict								=  ModelForTest()
   _snake_case							:		List[Any]								=  ModelHook()

   add_hook_to_module(a_,				a_		)
   add_hook_to_module(a_,				a_,				append=a_		)

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


 def 			UpperCamelCase_       (    self:		Tuple		):




   '''simple docstring'''



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

   _snake_case							:		Any								=  PreForwardHook()
   add_hook_to_module(a_,				a_		)
   _snake_case							:		List[Any]								=  test_model(a_		)
   self.assertTrue(torch.allclose(a_,				a_,				atol=1E-5		)		)

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

   # You need to use the sequential hook to chain two or more hooks
   _snake_case							:		str								=  SequentialHook(PreForwardHook(),				PreForwardHook()		)
   add_hook_to_module(a_,				a_		)

   _snake_case							:		str								=  test_model(a_		)
   assert torch.allclose(a_,				a_,				atol=1E-5		)


 def 			UpperCamelCase_       (    self:		Optional[int]		):




   '''simple docstring'''



   _snake_case							:		Optional[int]								=  ModelForTest()
   _snake_case							:		List[Any]								=  torch.randn(2,				3		)
   _snake_case							:		List[str]								=  test_model(a_		)

   _snake_case							:		List[Any]								=  PostForwardHook()
   add_hook_to_module(a_,				a_		)
   _snake_case							:		Union[str, Any]								=  test_model(a_		)
   self.assertTrue(torch.allclose(a_,				output + 1,				atol=1E-5		)		)

   # Attaching a hook to a model when it already has one replaces, does not chain
   _snake_case							:		Tuple								=  PostForwardHook()
   add_hook_to_module(a_,				a_		)
   _snake_case							:		Optional[Any]								=  test_model(a_		)
   self.assertTrue(torch.allclose(a_,				output + 1,				atol=1E-5		)		)

   # You need to use the sequential hook to chain two or more hooks
   _snake_case							:		Dict								=  SequentialHook(PostForwardHook(),				PostForwardHook()		)
   add_hook_to_module(a_,				a_		)

   _snake_case							:		List[str]								=  test_model(a_		)
   assert torch.allclose(a_,				output + 2,				atol=1E-5		)


 def 			UpperCamelCase_       (    self:		Optional[Any]		):




   '''simple docstring'''



   _snake_case							:		str								=  ModelForTest()
   _snake_case							:		Any								=  torch.randn(2,				3		)
   _snake_case							:		List[str]								=  test_model(a_		)

   _snake_case							:		List[Any]								=  PostForwardHook()
   add_hook_to_module(a_,				a_		)
   _snake_case							:		Dict								=  test_model(a_		)
   self.assertTrue(torch.allclose(a_,				output + 1		)		)
   self.assertTrue(outputa.requires_grad		)

   _snake_case							:		Union[str, Any]								=  True
   _snake_case							:		Dict								=  test_model(a_		)
   self.assertFalse(outputa.requires_grad		)


 @require_multi_gpu
 def 			UpperCamelCase_       (    self:		Tuple		):




   '''simple docstring'''



   _snake_case							:		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
   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
   _snake_case							:		Any								=  torch.randn(2,				3		)
   _snake_case							:		Any								=  model(a_		)
   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(a_,				AlignDevicesHook(io_same_device=a_		)		)
   _snake_case							:		int								=  torch.randn(2,				3		).to(0		)
   _snake_case							:		Optional[Any]								=  model(a_		)
   self.assertEqual(output.device,				torch.device(0		)		)


 def 			UpperCamelCase_       (    self:		Union[str, Any]		):




   '''simple docstring'''



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

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

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

   _snake_case							:		str								=  torch.randn(2,				3		)
   _snake_case							:		Optional[int]								=  model(a_		)
   self.assertEqual(output.device,				a_		)

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

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

   # 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"""		)		)

   _snake_case							:		List[Any]								=  torch.randn(2,				3		)
   _snake_case							:		Any								=  model(a_		)
   self.assertEqual(output.device,				a_		)

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




   '''simple docstring'''



   _snake_case							:		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
   _snake_case							:		Tuple								=  0 if torch.cuda.is_available() else """cpu"""
   attach_align_device_hook(a_,				execution_device=a_,				offload=a_		)

   # 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
   _snake_case							:		Optional[int]								=  torch.device(a_		)
   self.assertEqual(model.batchnorm.running_mean.device,				a_		)

   _snake_case							:		List[Any]								=  torch.randn(2,				3		)
   _snake_case							:		List[Any]								=  model(a_		)
   self.assertEqual(output.device,				a_		)

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

   # 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"""		)		)

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

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




   '''simple docstring'''



   _snake_case							:		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
   _snake_case							:		Any								=  0 if torch.cuda.is_available() else """cpu"""
   attach_align_device_hook(
       a_,				execution_device=a_,				offload=a_,				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
   _snake_case							:		Optional[Any]								=  torch.device(a_		)
   self.assertEqual(model.batchnorm.running_mean.device,				a_		)

   _snake_case							:		int								=  torch.randn(2,				3		)
   _snake_case							:		str								=  model(a_		)
   self.assertEqual(output.device,				a_		)

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

   # 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"""		)		)

   _snake_case							:		Optional[int]								=  torch.randn(2,				3		)
   _snake_case							:		Any								=  model(a_		)
   self.assertEqual(output.device,				a_		)

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


