



'''simple docstring'''



import gc
import os
import tempfile
import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNetaDConditionModel
from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
from diffusers.models.attention_processor import (
    Attention,
    AttnProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRAAttnProcessor2_0,
    LoRAXFormersAttnProcessor,
    XFormersAttnProcessor,
)
from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device


def 							__UpperCAmelCase				(				UpperCamelCase__				:nn.Module			)		->						Tuple:
				snake_case__      :   Dict      =    {}
				for name in unet.attn_processors.keys():
								snake_case__      :   Tuple      =    None if name.endswith('''attn1.processor'''			) else unet.config.cross_attention_dim
								if name.startswith('''mid_block'''			):
												snake_case__      :   Union[str, Any]      =    unet.config.block_out_channels[-1]
								elif name.startswith('''up_blocks'''			):
												snake_case__      :   Dict      =    int(name[len('''up_blocks.'''			)]			)
												snake_case__      :   int      =    list(reversed(unet.config.block_out_channels			)			)[block_id]
								elif name.startswith('''down_blocks'''			):
												snake_case__      :   Any      =    int(name[len('''down_blocks.'''			)]			)
												snake_case__      :   List[str]      =    unet.config.block_out_channels[block_id]
								snake_case__      :   Dict      =    (
								    LoRAAttnProcessor2_0 if hasattr(UpperCamelCase__   ,				'''scaled_dot_product_attention'''			) else LoRAAttnProcessor
								)
								snake_case__      :   Optional[Any]      =    lora_attn_processor_class(
								    hidden_size=UpperCamelCase__   ,				cross_attention_dim=UpperCamelCase__			)
				snake_case__      :   Dict      =    AttnProcsLayers(UpperCamelCase__			)
				return lora_attn_procs, unet_lora_layers


def 							__UpperCAmelCase				(				UpperCamelCase__				:nn.Module			)		->						str:
				snake_case__      :   Any      =    {}
				snake_case__      :   str      =    (
				    LoRAAttnProcessor2_0 if hasattr(UpperCamelCase__   ,				'''scaled_dot_product_attention'''			) else LoRAAttnProcessor
				)
				for name, module in text_encoder.named_modules():
								if name.endswith(UpperCamelCase__			):
												snake_case__      :   Dict      =    lora_attn_processor_class(
												    hidden_size=module.out_proj.out_features   ,				cross_attention_dim=UpperCamelCase__			)
				return text_lora_attn_procs


def 							__UpperCAmelCase				(				UpperCamelCase__				:nn.Module			)		->						int:
				snake_case__      :   int      =    create_text_encoder_lora_attn_procs(UpperCamelCase__			)
				snake_case__      :   Tuple      =    AttnProcsLayers(UpperCamelCase__			)
				return text_encoder_lora_layers


def 							__UpperCAmelCase				(				UpperCamelCase__				:List[str]   ,				UpperCamelCase__				:List[str]=False			)		->						Optional[int]:
				for _, attn_proc in text_lora_attn_procs.items():
								# set up.weights
								for layer_name, layer_module in attn_proc.named_modules():
												if layer_name.endswith('''_lora'''			):
																snake_case__      :   List[str]      =    (
																    torch.randn_like(layer_module.up.weight			)
																    if randn_weight
																    else torch.zeros_like(layer_module.up.weight			)
																)
																snake_case__      :   str      =    torch.nn.Parameter(UpperCamelCase__			)






class 			_SCREAMING_SNAKE_CASE				(unittest.TestCase      ):



			def 				lowerCAmelCase (     self		:						Optional[int]			)      ->      str:
							"""simple docstring"""





							torch.manual_seed(0			)
							snake_case__      :   Dict      =    UNetaDConditionModel(
							    block_out_channels=(32, 64)    ,			layers_per_block=2    ,			sample_size=32    ,			in_channels=4    ,			out_channels=4    ,			down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')    ,			up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')    ,			cross_attention_dim=32    ,			)
							snake_case__      :   Union[str, Any]      =    DDIMScheduler(
							    beta_start=0.00085    ,			beta_end=0.012    ,			beta_schedule='''scaled_linear'''    ,			clip_sample=__UpperCamelCase    ,			set_alpha_to_one=__UpperCamelCase    ,			steps_offset=1    ,			)
							torch.manual_seed(0			)
							snake_case__      :   str      =    AutoencoderKL(
							    block_out_channels=[32, 64]    ,			in_channels=3    ,			out_channels=3    ,			down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']    ,			up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']    ,			latent_channels=4    ,			)
							snake_case__      :   List[Any]      =    CLIPTextConfig(
							    bos_token_id=0    ,			eos_token_id=2    ,			hidden_size=32    ,			intermediate_size=37    ,			layer_norm_eps=1e-05    ,			num_attention_heads=4    ,			num_hidden_layers=5    ,			pad_token_id=1    ,			vocab_size=1000    ,			)
							snake_case__      :   Optional[Any]      =    CLIPTextModel(__UpperCamelCase			)
							snake_case__      :   Optional[Any]      =    CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''			)

							snake_case__  ,      snake_case__      :   Union[str, Any]      =    create_unet_lora_layers(__UpperCamelCase			)
							snake_case__      :   Optional[int]      =    create_text_encoder_lora_layers(__UpperCamelCase			)

							snake_case__      :   Tuple      =    {
							    '''unet''': unet,
							    '''scheduler''': scheduler,
							    '''vae''': vae,
							    '''text_encoder''': text_encoder,
							    '''tokenizer''': tokenizer,
							    '''safety_checker''': None,
							    '''feature_extractor''': None,
							}
							snake_case__      :   Optional[Any]      =    {
							    '''unet_lora_layers''': unet_lora_layers,
							    '''text_encoder_lora_layers''': text_encoder_lora_layers,
							    '''unet_lora_attn_procs''': unet_lora_attn_procs,
							}
							return pipeline_components, lora_components



			def 				lowerCAmelCase (     self		:						List[str]			)      ->      List[str]:
							"""simple docstring"""





							snake_case__      :   List[str]      =    1
							snake_case__      :   List[str]      =    10
							snake_case__      :   Dict      =    4
							snake_case__      :   Dict      =    (32, 32)

							snake_case__      :   Tuple      =    torch.manual_seed(0			)
							snake_case__      :   Tuple      =    floats_tensor((batch_size, num_channels) + sizes			)
							snake_case__      :   Any      =    torch.randint(1    ,			__UpperCamelCase    ,			size=(batch_size, sequence_length)    ,			generator=__UpperCamelCase			)

							snake_case__      :   List[str]      =    {
							    '''prompt''': '''A painting of a squirrel eating a burger''',
							    '''generator''': generator,
							    '''num_inference_steps''': 2,
							    '''guidance_scale''': 6.0,
							    '''output_type''': '''numpy''',
							}

							return noise, input_ids, pipeline_inputs

							# copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb



			def 				lowerCAmelCase (     self		:						Any			)      ->      List[str]:
							"""simple docstring"""





							snake_case__      :   List[Any]      =    77

							snake_case__      :   Optional[Any]      =    torch.randint(2    ,			56    ,			size=(1, max_seq_length)    ,			generator=torch.manual_seed(0			)			)

							snake_case__      :   Any      =    {}
							snake_case__      :   Any      =    inputs
							return prepared_inputs



			def 				lowerCAmelCase (     self		:						Tuple    ,			__UpperCamelCase		:						Tuple			)      ->      Union[str, Any]:
							"""simple docstring"""





							snake_case__  ,      snake_case__      :   Any      =    self.get_dummy_components()
							LoraLoaderMixin.save_lora_weights(
							    save_directory=__UpperCamelCase    ,			unet_lora_layers=lora_components['''unet_lora_layers''']    ,			text_encoder_lora_layers=lora_components['''text_encoder_lora_layers''']    ,			)
							self.assertTrue(os.path.isfile(os.path.join(__UpperCamelCase    ,			'''pytorch_lora_weights.bin'''			)			)			)



			def 				lowerCAmelCase (     self		:						Optional[int]			)      ->      List[Any]:
							"""simple docstring"""





							snake_case__  ,      snake_case__      :   Dict      =    self.get_dummy_components()
							snake_case__      :   List[Any]      =    StableDiffusionPipeline(**__UpperCamelCase			)
							snake_case__      :   str      =    sd_pipe.to(__UpperCamelCase			)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							snake_case__  ,      snake_case__  ,      snake_case__      :   str      =    self.get_dummy_inputs()

							snake_case__      :   Union[str, Any]      =    sd_pipe(**__UpperCamelCase			).images
							snake_case__      :   List[Any]      =    original_images[0, -3:, -3:, -1]

							with tempfile.TemporaryDirectory() as tmpdirname:
											LoraLoaderMixin.save_lora_weights(
											    save_directory=__UpperCamelCase    ,			unet_lora_layers=lora_components['''unet_lora_layers''']    ,			text_encoder_lora_layers=lora_components['''text_encoder_lora_layers''']    ,			)
											self.assertTrue(os.path.isfile(os.path.join(__UpperCamelCase    ,			'''pytorch_lora_weights.bin'''			)			)			)
											sd_pipe.load_lora_weights(__UpperCamelCase			)

							snake_case__      :   Optional[Any]      =    sd_pipe(**__UpperCamelCase			).images
							snake_case__      :   int      =    lora_images[0, -3:, -3:, -1]

							# Outputs shouldn't match.
							self.assertFalse(torch.allclose(torch.from_numpy(__UpperCamelCase			)    ,			torch.from_numpy(__UpperCamelCase			)			)			)



			def 				lowerCAmelCase (     self		:						List[Any]			)      ->      Optional[Any]:
							"""simple docstring"""





							snake_case__  ,      snake_case__      :   Dict      =    self.get_dummy_components()
							snake_case__      :   Union[str, Any]      =    StableDiffusionPipeline(**__UpperCamelCase			)
							snake_case__      :   List[Any]      =    sd_pipe.to(__UpperCamelCase			)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							snake_case__  ,      snake_case__  ,      snake_case__      :   str      =    self.get_dummy_inputs()

							snake_case__      :   List[str]      =    sd_pipe(**__UpperCamelCase			).images
							snake_case__      :   Optional[Any]      =    original_images[0, -3:, -3:, -1]

							with tempfile.TemporaryDirectory() as tmpdirname:
											LoraLoaderMixin.save_lora_weights(
											    save_directory=__UpperCamelCase    ,			unet_lora_layers=lora_components['''unet_lora_layers''']    ,			text_encoder_lora_layers=lora_components['''text_encoder_lora_layers''']    ,			safe_serialization=__UpperCamelCase    ,			)
											self.assertTrue(os.path.isfile(os.path.join(__UpperCamelCase    ,			'''pytorch_lora_weights.safetensors'''			)			)			)
											sd_pipe.load_lora_weights(__UpperCamelCase			)

							snake_case__      :   List[str]      =    sd_pipe(**__UpperCamelCase			).images
							snake_case__      :   Dict      =    lora_images[0, -3:, -3:, -1]

							# Outputs shouldn't match.
							self.assertFalse(torch.allclose(torch.from_numpy(__UpperCamelCase			)    ,			torch.from_numpy(__UpperCamelCase			)			)			)



			def 				lowerCAmelCase (     self		:						str			)      ->      int:
							"""simple docstring"""





							snake_case__  ,      snake_case__      :   str      =    self.get_dummy_components()
							snake_case__      :   Any      =    lora_components['''unet_lora_attn_procs''']
							snake_case__      :   Tuple      =    StableDiffusionPipeline(**__UpperCamelCase			)
							snake_case__      :   Dict      =    sd_pipe.to(__UpperCamelCase			)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							snake_case__  ,      snake_case__  ,      snake_case__      :   List[Any]      =    self.get_dummy_inputs()

							snake_case__      :   List[str]      =    sd_pipe(**__UpperCamelCase			).images
							snake_case__      :   Optional[Any]      =    original_images[0, -3:, -3:, -1]

							with tempfile.TemporaryDirectory() as tmpdirname:
											snake_case__      :   Dict      =    sd_pipe.unet
											unet.set_attn_processor(__UpperCamelCase			)
											unet.save_attn_procs(__UpperCamelCase			)
											self.assertTrue(os.path.isfile(os.path.join(__UpperCamelCase    ,			'''pytorch_lora_weights.bin'''			)			)			)
											sd_pipe.load_lora_weights(__UpperCamelCase			)

							snake_case__      :   Tuple      =    sd_pipe(**__UpperCamelCase			).images
							snake_case__      :   str      =    lora_images[0, -3:, -3:, -1]

							# Outputs shouldn't match.
							self.assertFalse(torch.allclose(torch.from_numpy(__UpperCamelCase			)    ,			torch.from_numpy(__UpperCamelCase			)			)			)



			def 				lowerCAmelCase (     self		:						Tuple			)      ->      Union[str, Any]:
							"""simple docstring"""





							snake_case__  ,      snake_case__      :   Any      =    self.get_dummy_components()
							snake_case__      :   List[Any]      =    StableDiffusionPipeline(**__UpperCamelCase			)

							snake_case__      :   Dict      =    self.get_dummy_tokens()

							# inference without lora
							snake_case__      :   Optional[int]      =    pipe.text_encoder(**__UpperCamelCase			)[0]
							assert outputs_without_lora.shape == (1, 77, 32)

							# create lora_attn_procs with zeroed out up.weights
							snake_case__      :   Optional[Any]      =    create_text_encoder_lora_attn_procs(pipe.text_encoder			)
							set_lora_up_weights(__UpperCamelCase    ,			randn_weight=__UpperCamelCase			)

							# monkey patch
							pipe._modify_text_encoder(__UpperCamelCase			)

							# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
							del text_attn_procs
							gc.collect()

							# inference with lora
							snake_case__      :   Dict      =    pipe.text_encoder(**__UpperCamelCase			)[0]
							assert outputs_with_lora.shape == (1, 77, 32)

							assert torch.allclose(
							    __UpperCamelCase    ,			__UpperCamelCase			), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"

							# create lora_attn_procs with randn up.weights
							snake_case__      :   int      =    create_text_encoder_lora_attn_procs(pipe.text_encoder			)
							set_lora_up_weights(__UpperCamelCase    ,			randn_weight=__UpperCamelCase			)

							# monkey patch
							pipe._modify_text_encoder(__UpperCamelCase			)

							# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
							del text_attn_procs
							gc.collect()

							# inference with lora
							snake_case__      :   Any      =    pipe.text_encoder(**__UpperCamelCase			)[0]
							assert outputs_with_lora.shape == (1, 77, 32)

							assert not torch.allclose(
							    __UpperCamelCase    ,			__UpperCamelCase			), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"



			def 				lowerCAmelCase (     self		:						List[str]			)      ->      int:
							"""simple docstring"""





							snake_case__  ,      snake_case__      :   int      =    self.get_dummy_components()
							snake_case__      :   Optional[int]      =    StableDiffusionPipeline(**__UpperCamelCase			)

							snake_case__      :   int      =    self.get_dummy_tokens()

							# inference without lora
							snake_case__      :   List[Any]      =    pipe.text_encoder(**__UpperCamelCase			)[0]
							assert outputs_without_lora.shape == (1, 77, 32)

							# create lora_attn_procs with randn up.weights
							snake_case__      :   Any      =    create_text_encoder_lora_attn_procs(pipe.text_encoder			)
							set_lora_up_weights(__UpperCamelCase    ,			randn_weight=__UpperCamelCase			)

							# monkey patch
							pipe._modify_text_encoder(__UpperCamelCase			)

							# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
							del text_attn_procs
							gc.collect()

							# inference with lora
							snake_case__      :   Tuple      =    pipe.text_encoder(**__UpperCamelCase			)[0]
							assert outputs_with_lora.shape == (1, 77, 32)

							assert not torch.allclose(
							    __UpperCamelCase    ,			__UpperCamelCase			), "lora outputs should be different to without lora outputs"

							# remove monkey patch
							pipe._remove_text_encoder_monkey_patch()

							# inference with removed lora
							snake_case__      :   List[Any]      =    pipe.text_encoder(**__UpperCamelCase			)[0]
							assert outputs_without_lora_removed.shape == (1, 77, 32)

							assert torch.allclose(
							    __UpperCamelCase    ,			__UpperCamelCase			), "remove lora monkey patch should restore the original outputs"



			def 				lowerCAmelCase (     self		:						Optional[Any]			)      ->      Tuple:
							"""simple docstring"""





							snake_case__  ,      snake_case__      :   Tuple      =    self.get_dummy_components()
							snake_case__      :   Tuple      =    StableDiffusionPipeline(**__UpperCamelCase			)
							snake_case__      :   Any      =    sd_pipe.to(__UpperCamelCase			)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							snake_case__  ,      snake_case__  ,      snake_case__      :   str      =    self.get_dummy_inputs()

							with tempfile.TemporaryDirectory() as tmpdirname:
											LoraLoaderMixin.save_lora_weights(
											    save_directory=__UpperCamelCase    ,			unet_lora_layers=lora_components['''unet_lora_layers''']    ,			text_encoder_lora_layers=lora_components['''text_encoder_lora_layers''']    ,			)
											self.assertTrue(os.path.isfile(os.path.join(__UpperCamelCase    ,			'''pytorch_lora_weights.bin'''			)			)			)
											sd_pipe.load_lora_weights(__UpperCamelCase			)

							snake_case__      :   Dict      =    sd_pipe(**__UpperCamelCase			).images
							snake_case__      :   Optional[Any]      =    lora_images[0, -3:, -3:, -1]

							snake_case__      :   Optional[Any]      =    sd_pipe(**__UpperCamelCase    ,			cross_attention_kwargs={'''scale''': 0.5}			).images
							snake_case__      :   Union[str, Any]      =    lora_images_with_scale[0, -3:, -3:, -1]

							# Outputs shouldn't match.
							self.assertFalse(
							    torch.allclose(torch.from_numpy(__UpperCamelCase			)    ,			torch.from_numpy(__UpperCamelCase			)			)			)



			def 				lowerCAmelCase (     self		:						Tuple			)      ->      Optional[Any]:
							"""simple docstring"""





							with tempfile.TemporaryDirectory() as tmpdirname:
											self.create_lora_weight_file(__UpperCamelCase			)

											snake_case__  ,      snake_case__      :   str      =    self.get_dummy_components()
											snake_case__      :   Optional[Any]      =    StableDiffusionPipeline(**__UpperCamelCase			)
											snake_case__      :   List[str]      =    sd_pipe.to(__UpperCamelCase			)
											sd_pipe.set_progress_bar_config(disable=__UpperCamelCase			)

											# check if vanilla attention processors are used
											for _, module in sd_pipe.unet.named_modules():
															if isinstance(__UpperCamelCase    ,			__UpperCamelCase			):
																			self.assertIsInstance(module.processor    ,			(AttnProcessor, AttnProcessor2_0)			)

            # load LoRA weight file
											sd_pipe.load_lora_weights(__UpperCamelCase			)

											# check if lora attention processors are used
											for _, module in sd_pipe.unet.named_modules():
															if isinstance(__UpperCamelCase    ,			__UpperCamelCase			):
																			snake_case__      :   Optional[Any]      =    (
																			    LoRAAttnProcessor2_0 if hasattr(__UpperCamelCase    ,			'''scaled_dot_product_attention'''			) else LoRAAttnProcessor
																			)
																			self.assertIsInstance(module.processor    ,			__UpperCamelCase			)



			@unittest.skipIf(torch_device != '''cuda'''    ,			'''This test is supposed to run on GPU'''			)
			def 				lowerCAmelCase (     self		:						List[str]			)      ->      Union[str, Any]:
							"""simple docstring"""





							with tempfile.TemporaryDirectory() as tmpdirname:
											self.create_lora_weight_file(__UpperCamelCase			)

											snake_case__  ,      snake_case__      :   int      =    self.get_dummy_components()
											snake_case__      :   int      =    StableDiffusionPipeline(**__UpperCamelCase			)
											snake_case__      :   Any      =    sd_pipe.to(__UpperCamelCase			)
											sd_pipe.set_progress_bar_config(disable=__UpperCamelCase			)

											# enable XFormers
											sd_pipe.enable_xformers_memory_efficient_attention()

											# check if xFormers attention processors are used
											for _, module in sd_pipe.unet.named_modules():
															if isinstance(__UpperCamelCase    ,			__UpperCamelCase			):
																			self.assertIsInstance(module.processor    ,			__UpperCamelCase			)

            # load LoRA weight file
											sd_pipe.load_lora_weights(__UpperCamelCase			)

											# check if lora attention processors are used
											for _, module in sd_pipe.unet.named_modules():
															if isinstance(__UpperCamelCase    ,			__UpperCamelCase			):
																			self.assertIsInstance(module.processor    ,			__UpperCamelCase			)







			@unittest.skipIf(torch_device != '''cuda'''    ,			'''This test is supposed to run on GPU'''			)
			def 				lowerCAmelCase (     self		:						Dict			)      ->      Union[str, Any]:
							"""simple docstring"""





							snake_case__  ,      snake_case__      :   Optional[Any]      =    self.get_dummy_components()
							snake_case__      :   List[Any]      =    StableDiffusionPipeline(**__UpperCamelCase			)
							snake_case__      :   Optional[int]      =    sd_pipe.to(__UpperCamelCase			)
							sd_pipe.set_progress_bar_config(disable=__UpperCamelCase			)

							snake_case__  ,      snake_case__  ,      snake_case__      :   Any      =    self.get_dummy_inputs()

							# enable XFormers
							sd_pipe.enable_xformers_memory_efficient_attention()

							snake_case__      :   Union[str, Any]      =    sd_pipe(**__UpperCamelCase			).images
							snake_case__      :   List[str]      =    original_images[0, -3:, -3:, -1]

							with tempfile.TemporaryDirectory() as tmpdirname:
											LoraLoaderMixin.save_lora_weights(
											    save_directory=__UpperCamelCase    ,			unet_lora_layers=lora_components['''unet_lora_layers''']    ,			text_encoder_lora_layers=lora_components['''text_encoder_lora_layers''']    ,			)
											self.assertTrue(os.path.isfile(os.path.join(__UpperCamelCase    ,			'''pytorch_lora_weights.bin'''			)			)			)
											sd_pipe.load_lora_weights(__UpperCamelCase			)

							snake_case__      :   str      =    sd_pipe(**__UpperCamelCase			).images
							snake_case__      :   Optional[int]      =    lora_images[0, -3:, -3:, -1]

							# Outputs shouldn't match.
							self.assertFalse(torch.allclose(torch.from_numpy(__UpperCamelCase			)    ,			torch.from_numpy(__UpperCamelCase			)			)			)









'''simple docstring'''



import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
from packaging import version
from safetensors.torch import storage_ptr, storage_size
from torch import nn

from .utils import logging


_lowercase       :						int						=[nn.LayerNorm]

_lowercase       :						Tuple						=logging.get_logger(__name__)

_lowercase       :						List[Any]						=version.parse(version.parse(torch.__version__).base_version)

_lowercase       :						int						=parsed_torch_version_base >= version.parse("2.0")
_lowercase       :						Tuple						=parsed_torch_version_base >= version.parse("1.12")
_lowercase       :						int						=parsed_torch_version_base >= version.parse("1.11")
_lowercase       :						Dict						=parsed_torch_version_base < version.parse("1.11")


def 							__UpperCAmelCase				(				UpperCamelCase__				:Dict   ,				UpperCamelCase__				:List[str]   ,				UpperCamelCase__				:str   ,				UpperCamelCase__				:Any   ,				self				:Optional[int]			)		->						Union[str, Any]:
				from torch import _softmax_backward_data

				if is_torch_less_than_1_11:
								return _softmax_backward_data(UpperCamelCase__   ,				UpperCamelCase__   ,				parent.dim   ,				self			)
				else:
								return _softmax_backward_data(UpperCamelCase__   ,				UpperCamelCase__   ,				parent.dim   ,				self.dtype			)


def 							__UpperCAmelCase				(				UpperCamelCase__				:nn.Linear   ,				UpperCamelCase__				:torch.LongTensor   ,				UpperCamelCase__				:int = 0			)		->						nn.Linear:
				snake_case__      :   Optional[Any]      =    index.to(layer.weight.device			)
				snake_case__      :   int      =    layer.weight.index_select(UpperCamelCase__   ,				UpperCamelCase__			).clone().detach()
				if layer.bias is not None:
								if dim == 1:
												snake_case__      :   Dict      =    layer.bias.clone().detach()
								else:
												snake_case__      :   Optional[Any]      =    layer.bias[index].clone().detach()
				snake_case__      :   Union[str, Any]      =    list(layer.weight.size()			)
				snake_case__      :   int      =    len(UpperCamelCase__			)
				snake_case__      :   int      =    nn.Linear(new_size[1]   ,				new_size[0]   ,				bias=layer.bias is not None			).to(layer.weight.device			)
				snake_case__      :   Union[str, Any]      =    False
				new_layer.weight.copy_(W.contiguous()			)
				snake_case__      :   List[Any]      =    True
				if layer.bias is not None:
								snake_case__      :   Union[str, Any]      =    False
								new_layer.bias.copy_(b.contiguous()			)
								snake_case__      :   Dict      =    True
				return new_layer






class 			_SCREAMING_SNAKE_CASE				(nn.Module      ):



			def __init__(     self		:						int    ,			__UpperCamelCase		:						Tuple    ,			__UpperCamelCase		:						Optional[Any]			)      ->      Union[str, Any]:
							"""simple docstring"""





							super().__init__()
							snake_case__      :   str      =    nf
							snake_case__      :   Tuple      =    nn.Parameter(torch.empty(__UpperCamelCase    ,			__UpperCamelCase			)			)
							snake_case__      :   Optional[Any]      =    nn.Parameter(torch.zeros(__UpperCamelCase			)			)
							nn.init.normal_(self.weight    ,			std=0.02			)







			def 				lowerCAmelCase (     self		:						Dict    ,			__UpperCamelCase		:						Optional[int]			)      ->      Dict:
							"""simple docstring"""





							snake_case__      :   Dict      =    x.size()[:-1] + (self.nf,)
							snake_case__      :   List[str]      =    torch.addmm(self.bias    ,			x.view(-1    ,			x.size(-1			)			)    ,			self.weight			)
							snake_case__      :   Dict      =    x.view(__UpperCamelCase			)
							return x



def 							__UpperCAmelCase				(				UpperCamelCase__				:ConvaD   ,				UpperCamelCase__				:torch.LongTensor   ,				UpperCamelCase__				:int = 1			)		->						ConvaD:
				snake_case__      :   Optional[Any]      =    index.to(layer.weight.device			)
				snake_case__      :   List[Any]      =    layer.weight.index_select(UpperCamelCase__   ,				UpperCamelCase__			).clone().detach()
				if dim == 0:
								snake_case__      :   Any      =    layer.bias.clone().detach()
				else:
								snake_case__      :   Optional[Any]      =    layer.bias[index].clone().detach()
				snake_case__      :   Union[str, Any]      =    list(layer.weight.size()			)
				snake_case__      :   List[str]      =    len(UpperCamelCase__			)
				snake_case__      :   Optional[Any]      =    ConvaD(new_size[1]   ,				new_size[0]			).to(layer.weight.device			)
				snake_case__      :   List[str]      =    False
				new_layer.weight.copy_(W.contiguous()			)
				snake_case__      :   List[Any]      =    True
				snake_case__      :   List[str]      =    False
				new_layer.bias.copy_(b.contiguous()			)
				snake_case__      :   str      =    True
				return new_layer


def 							__UpperCAmelCase				(				UpperCamelCase__				:Union[nn.Linear, ConvaD]   ,				UpperCamelCase__				:torch.LongTensor   ,				UpperCamelCase__				:Optional[int] = None			)		->						Union[nn.Linear, ConvaD]:
				if isinstance(UpperCamelCase__   ,				nn.Linear			):
								return prune_linear_layer(UpperCamelCase__   ,				UpperCamelCase__   ,				dim=0 if dim is None else dim			)
				elif isinstance(UpperCamelCase__   ,				UpperCamelCase__			):
								return prune_convad_layer(UpperCamelCase__   ,				UpperCamelCase__   ,				dim=1 if dim is None else dim			)
				else:
								raise ValueError(F'''Can\'t prune layer of class {layer.__class__}'''			)


def 							__UpperCAmelCase				(				UpperCamelCase__				:Callable[..., torch.Tensor]   ,				UpperCamelCase__				:int   ,				UpperCamelCase__				:int   ,				*UpperCamelCase__				:List[str]			)		->						torch.Tensor:
				assert len(UpperCamelCase__			) > 0, F'''{input_tensors} has to be a tuple/list of tensors'''

				# inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
				snake_case__      :   str      =    len(inspect.signature(UpperCamelCase__			).parameters			)
				if num_args_in_forward_chunk_fn != len(UpperCamelCase__			):
								raise ValueError(
								    F'''forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(UpperCamelCase__			)} input '''
								    '''tensors are given'''			)

				if chunk_size > 0:
								snake_case__      :   Tuple      =    input_tensors[0].shape[chunk_dim]
								for input_tensor in input_tensors:
												if input_tensor.shape[chunk_dim] != tensor_shape:
																raise ValueError(
																    F'''All input tenors have to be of the same shape: {tensor_shape}, '''
																    F'''found shape {input_tensor.shape[chunk_dim]}'''			)

								if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
												raise ValueError(
												    F'''The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk '''
												    F'''size {chunk_size}'''			)

								snake_case__      :   Optional[Any]      =    input_tensors[0].shape[chunk_dim] // chunk_size

								# chunk input tensor into tuples
								snake_case__      :   List[Any]      =    tuple(input_tensor.chunk(UpperCamelCase__   ,				dim=UpperCamelCase__			) for input_tensor in input_tensors			)
								# apply forward fn to every tuple
								snake_case__      :   List[Any]      =    tuple(forward_fn(*UpperCamelCase__			) for input_tensors_chunk in zip(*UpperCamelCase__			)			)
								# concatenate output at same dimension
								return torch.cat(UpperCamelCase__   ,				dim=UpperCamelCase__			)

				return forward_fn(*UpperCamelCase__			)


def 							__UpperCAmelCase				(				UpperCamelCase__				:List[int]   ,				UpperCamelCase__				:int   ,				UpperCamelCase__				:int   ,				UpperCamelCase__				:Set[int]			)		->						Tuple[Set[int], torch.LongTensor]:
				snake_case__      :   str      =    torch.ones(UpperCamelCase__   ,				UpperCamelCase__			)
				snake_case__      :   int      =    set(UpperCamelCase__			) - already_pruned_heads  # Convert to set and remove already pruned heads
				for head in heads:
								# Compute how many pruned heads are before the head and move the index accordingly
								snake_case__      :   str      =    head - sum(1 if h < head else 0 for h in already_pruned_heads			)
								snake_case__      :   Optional[int]      =    0
				snake_case__      :   Optional[Any]      =    mask.view(-1			).contiguous().eq(1			)
				snake_case__      :   torch.LongTensor      =    torch.arange(len(UpperCamelCase__			)			)[mask].long()
				return heads, index


def 							__UpperCAmelCase				(				*UpperCamelCase__				:Union[torch.Tensor, List[torch.Tensor]]   ,				UpperCamelCase__				:Optional[str] = None			)		->						Tuple[torch.Tensor, ...]:
				return torch.meshgrid(*UpperCamelCase__   ,				indexing=UpperCamelCase__			)



def 							__UpperCAmelCase				(				UpperCamelCase__				:torch.Tensor			)		->						Tuple[torch.device, int, int]:
				return tensor.device, storage_ptr(UpperCamelCase__			), storage_size(UpperCamelCase__			)




