


'''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


_lowerCAmelCase	:		Optional[Any]		  =						[nn.LayerNorm]

_lowerCAmelCase	:		Optional[Any]		  =						logging.get_logger(__name__)

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

_lowerCAmelCase	:		List[str]		  =						parsed_torch_version_base >= version.parse('2.0')
_lowerCAmelCase	:		List[Any]		  =						parsed_torch_version_base >= version.parse('1.12')
_lowerCAmelCase	:		Union[str, Any]		  =						parsed_torch_version_base >= version.parse('1.11')
_lowerCAmelCase	:		Dict		  =						parsed_torch_version_base < version.parse('1.11')

def       __UpperCamelCase     (     _A     :					Union[str, Any]       , _A     :					Any       , _A     :					Union[str, Any]       , _A     :					Tuple       , self     :					Union[str, Any]					)     ->   int:




	"""simple docstring"""





	from torch import _softmax_backward_data

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

def       __UpperCamelCase     (     _A     :					nn.Linear       , _A     :					torch.LongTensor       , _A     :					int = 0					)     ->   nn.Linear:




	"""simple docstring"""





	lowerCAmelCase       :     Union[str, Any]   				=			index.to(layer.weight.device					)
	lowerCAmelCase       :     Any   				=			layer.weight.index_select(_A       , _A					).clone().detach()
	if layer.bias is not None:
		if dim == 1:
			lowerCAmelCase       :     str   				=			layer.bias.clone().detach()
		else:
			lowerCAmelCase       :     Optional[Any]   				=			layer.bias[index].clone().detach()
	lowerCAmelCase       :     int   				=			list(layer.weight.size()					)
	lowerCAmelCase       :     Union[str, Any]   				=			len(_A					)
	lowerCAmelCase       :     int   				=			nn.Linear(new_size[1]       , new_size[0]       , bias=layer.bias is not None					).to(layer.weight.device					)
	lowerCAmelCase       :     Optional[Any]   				=			False
	new_layer.weight.copy_(W.contiguous()					)
	lowerCAmelCase       :     Union[str, Any]   				=			True
	if layer.bias is not None:
		lowerCAmelCase       :     Any   				=			False
		new_layer.bias.copy_(b.contiguous()					)
		lowerCAmelCase       :     int   				=			True
	return new_layer



class 			lowerCAmelCase					(					nn.Module				):



	def __init__(       self			,       snake_case__			,       snake_case__     ):
		super().__init__()
		lowerCAmelCase       :     Optional[Any]   				=			nf
		lowerCAmelCase       :     Optional[Any]   				=			nn.Parameter(torch.empty(snake_case__			,       snake_case__     )     )
		lowerCAmelCase       :     Any   				=			nn.Parameter(torch.zeros(snake_case__     )     )
		nn.init.normal_(self.weight			,       std=0.0_2     )



	def 					lowercase  (       self			,       snake_case__     ):
		lowerCAmelCase       :     List[Any]   				=			x.size()[:-1] + (self.nf,)
		lowerCAmelCase       :     Any   				=			torch.addmm(self.bias			,       x.view(-1			,       x.size(-1     )     )			,       self.weight     )
		lowerCAmelCase       :     List[Any]   				=			x.view(snake_case__     )
		return x


def       __UpperCamelCase     (     _A     :					ConvaD       , _A     :					torch.LongTensor       , _A     :					int = 1					)     ->   ConvaD:




	"""simple docstring"""





	lowerCAmelCase       :     List[str]   				=			index.to(layer.weight.device					)
	lowerCAmelCase       :     Dict   				=			layer.weight.index_select(_A       , _A					).clone().detach()
	if dim == 0:
		lowerCAmelCase       :     str   				=			layer.bias.clone().detach()
	else:
		lowerCAmelCase       :     Union[str, Any]   				=			layer.bias[index].clone().detach()
	lowerCAmelCase       :     int   				=			list(layer.weight.size()					)
	lowerCAmelCase       :     List[str]   				=			len(_A					)
	lowerCAmelCase       :     Tuple   				=			ConvaD(new_size[1]       , new_size[0]					).to(layer.weight.device					)
	lowerCAmelCase       :     Any   				=			False
	new_layer.weight.copy_(W.contiguous()					)
	lowerCAmelCase       :     Union[str, Any]   				=			True
	lowerCAmelCase       :     Optional[Any]   				=			False
	new_layer.bias.copy_(b.contiguous()					)
	lowerCAmelCase       :     Optional[Any]   				=			True
	return new_layer

def       __UpperCamelCase     (     _A     :					Union[nn.Linear, ConvaD]       , _A     :					torch.LongTensor       , _A     :					Optional[int] = None					)     ->   Union[nn.Linear, ConvaD]:




	"""simple docstring"""





	if isinstance(_A       , nn.Linear					):
		return prune_linear_layer(_A       , _A       , dim=0 if dim is None else dim					)
	elif isinstance(_A       , _A					):
		return prune_convad_layer(_A       , _A       , dim=1 if dim is None else dim					)
	else:
		raise ValueError(F"Can't prune layer of class {layer.__class__}"					)

def       __UpperCamelCase     (     _A     :					Callable[..., torch.Tensor]       , _A     :					int       , _A     :					int       , *_A     :					int					)     ->   torch.Tensor:




	"""simple docstring"""





	assert len(_A					) > 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
	lowerCAmelCase       :     Optional[Any]   				=			len(inspect.signature(_A					).parameters					)
	if num_args_in_forward_chunk_fn != len(_A					):
		raise ValueError(
		    F"forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(_A					)} input "
		    'tensors are given'					)

	if chunk_size > 0:
		lowerCAmelCase       :     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}"					)

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

		# chunk input tensor into tuples
		lowerCAmelCase       :     Dict   				=			tuple(input_tensor.chunk(_A       , dim=_A					) for input_tensor in input_tensors					)
		# apply forward fn to every tuple
		lowerCAmelCase       :     str   				=			tuple(forward_fn(*_A					) for input_tensors_chunk in zip(*_A					)					)
		# concatenate output at same dimension
		return torch.cat(_A       , dim=_A					)

	return forward_fn(*_A					)

def       __UpperCamelCase     (     _A     :					List[int]       , _A     :					int       , _A     :					int       , _A     :					Set[int]					)     ->   Tuple[Set[int], torch.LongTensor]:




	"""simple docstring"""





	lowerCAmelCase       :     Any   				=			torch.ones(_A       , _A					)
	lowerCAmelCase       :     Dict   				=			set(_A					) - 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
		lowerCAmelCase       :     List[Any]   				=			head - sum(1 if h < head else 0 for h in already_pruned_heads					)
		lowerCAmelCase       :     Optional[int]   				=			0
	lowerCAmelCase       :     int   				=			mask.view(-1					).contiguous().eq(1					)
	lowerCAmelCase       :     torch.LongTensor   				=			torch.arange(len(_A					)					)[mask].long()
	return heads, index

def       __UpperCamelCase     (     *_A     :					Union[torch.Tensor, List[torch.Tensor]]       , _A     :					Optional[str] = None					)     ->   Tuple[torch.Tensor, ...]:




	"""simple docstring"""





	return torch.meshgrid(*_A       , indexing=_A					)






def       __UpperCamelCase     (     _A     :					torch.Tensor					)     ->   Tuple[torch.device, int, int]:




	"""simple docstring"""





	return tensor.device, storage_ptr(_A					), storage_size(_A					)




'''simple docstring'''




from collections import OrderedDict
from dataclasses import fields
from typing import Any, Tuple

import numpy as np

from .import_utils import is_torch_available

def       __UpperCamelCase     (     _A     :					List[str]					)     ->   Any:




	"""simple docstring"""





	if is_torch_available():
		import torch

		if isinstance(_A       , torch.Tensor					):
			return True

	return isinstance(_A       , np.ndarray					)






class 			lowerCAmelCase					(					a				):



	def 					lowercase  (       self     ):
		lowerCAmelCase       :     List[str]   				=			fields(self     )

		# Safety and consistency checks
		if not len(snake_case__     ):
			raise ValueError(f"{self.__class__.__name__} has no fields."     )

		lowerCAmelCase       :     List[str]   				=			getattr(self			,       class_fields[0].name     )
		lowerCAmelCase       :     List[Any]   				=			all(getattr(self			,       field.name     ) is None for field in class_fields[1:]     )

		if other_fields_are_none and isinstance(snake_case__			,       snake_case__     ):
			for key, value in first_field.items():
				lowerCAmelCase       :     Optional[int]   				=			value
		else:
			for field in class_fields:
				lowerCAmelCase       :     str   				=			getattr(self			,       field.name     )
				if v is not None:
					lowerCAmelCase       :     Optional[Any]   				=			v



	def __delitem__(       self			,       *snake_case__			,       **snake_case__     ):
		raise Exception(f"You cannot use ``__delitem__`` on a {self.__class__.__name__} instance."     )



	def 					lowercase  (       self			,       *snake_case__			,       **snake_case__     ):
		raise Exception(f"You cannot use ``setdefault`` on a {self.__class__.__name__} instance."     )



	def 					lowercase  (       self			,       *snake_case__			,       **snake_case__     ):
		raise Exception(f"You cannot use ``pop`` on a {self.__class__.__name__} instance."     )



	def 					lowercase  (       self			,       *snake_case__			,       **snake_case__     ):
		raise Exception(f"You cannot use ``update`` on a {self.__class__.__name__} instance."     )



	def __getitem__(       self			,       snake_case__     ):
		if isinstance(snake_case__			,       snake_case__     ):
			lowerCAmelCase       :     Dict   				=			dict(self.items()     )
			return inner_dict[k]
		else:
			return self.to_tuple()[k]



	def __setattr__(       self			,       snake_case__			,       snake_case__     ):
		if name in self.keys() and value is not None:
			# Don't call self.__setitem__ to avoid recursion errors
			super().__setitem__(snake_case__			,       snake_case__     )
		super().__setattr__(snake_case__			,       snake_case__     )



	def __setitem__(       self			,       snake_case__			,       snake_case__     ):
		# Will raise a KeyException if needed
		super().__setitem__(snake_case__			,       snake_case__     )
		# Don't call self.__setattr__ to avoid recursion errors
		super().__setattr__(snake_case__			,       snake_case__     )



	def 					lowercase  (       self     ):
		return tuple(self[k] for k in self.keys()     )
