





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


UpperCAmelCase__    :      List[Any]       					=[nn.LayerNorm]

UpperCAmelCase__    :      Dict       					=logging.get_logger(__name__)

UpperCAmelCase__    :      Dict       					=version.parse(version.parse(torch.__version__).base_version)

UpperCAmelCase__    :      Optional[Any]       					=parsed_torch_version_base >= version.parse('''2.0''')
UpperCAmelCase__    :      str       					=parsed_torch_version_base >= version.parse('''1.12''')
UpperCAmelCase__    :      Dict       					=parsed_torch_version_base >= version.parse('''1.11''')
UpperCAmelCase__    :      Tuple       					=parsed_torch_version_base < version.parse('''1.11''')





def   _lowercase       (    _UpperCAmelCase      ,   _UpperCAmelCase      ,   _UpperCAmelCase      ,   _UpperCAmelCase      ,   self					)    ->	int:
				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   _lowercase       (    _UpperCAmelCase      ,   _UpperCAmelCase      ,   _UpperCAmelCase = 0					)    ->	nn.Linear:
				lowerCamelCase									=index.to(layer.weight.device					)
				lowerCamelCase									=layer.weight.index_select(_UpperCAmelCase      ,   _UpperCAmelCase					).clone().detach()
				if layer.bias is not None:
								if dim == 1:
												lowerCamelCase									=layer.bias.clone().detach()
								else:
												lowerCamelCase									=layer.bias[index].clone().detach()
				lowerCamelCase									=list(layer.weight.size()					)
				lowerCamelCase									=len(_UpperCAmelCase					)
				lowerCamelCase									=nn.Linear(new_size[1]      ,   new_size[0]      ,   bias=layer.bias is not None					).to(layer.weight.device					)
				lowerCamelCase									=False
				new_layer.weight.copy_(W.contiguous()					)
				lowerCamelCase									=True
				if layer.bias is not None:
								lowerCamelCase									=False
								new_layer.bias.copy_(b.contiguous()					)
								lowerCamelCase									=True
				return new_layer

class    __A  (    nn.Module     ):


				def __init__(		self ,			UpperCAmelCase_ ,			UpperCAmelCase_				):
								super().__init__()
								lowerCamelCase									=nf
								lowerCamelCase									=nn.Parameter(torch.empty(UpperCAmelCase_ ,			UpperCAmelCase_				)				)
								lowerCamelCase									=nn.Parameter(torch.zeros(UpperCAmelCase_				)				)
								nn.init.normal_(self.weight ,			std=0.0_2				)


				def _snake_case							(		self ,			UpperCAmelCase_				):
								lowerCamelCase									=x.size()[:-1] + (self.nf,)
								lowerCamelCase									=torch.addmm(self.bias ,			x.view(-1 ,			x.size(-1				)				) ,			self.weight				)
								lowerCamelCase									=x.view(UpperCAmelCase_				)
								return x





def   _lowercase       (    _UpperCAmelCase      ,   _UpperCAmelCase      ,   _UpperCAmelCase = 1					)    ->	ConvaD:
				lowerCamelCase									=index.to(layer.weight.device					)
				lowerCamelCase									=layer.weight.index_select(_UpperCAmelCase      ,   _UpperCAmelCase					).clone().detach()
				if dim == 0:
								lowerCamelCase									=layer.bias.clone().detach()
				else:
								lowerCamelCase									=layer.bias[index].clone().detach()
				lowerCamelCase									=list(layer.weight.size()					)
				lowerCamelCase									=len(_UpperCAmelCase					)
				lowerCamelCase									=ConvaD(new_size[1]      ,   new_size[0]					).to(layer.weight.device					)
				lowerCamelCase									=False
				new_layer.weight.copy_(W.contiguous()					)
				lowerCamelCase									=True
				lowerCamelCase									=False
				new_layer.bias.copy_(b.contiguous()					)
				lowerCamelCase									=True
				return new_layer





def   _lowercase       (    _UpperCAmelCase      ,   _UpperCAmelCase      ,   _UpperCAmelCase = 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   _lowercase       (    _UpperCAmelCase      ,   _UpperCAmelCase      ,   _UpperCAmelCase      ,   *_UpperCAmelCase					)    ->	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
				lowerCamelCase									=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:
								lowerCamelCase									=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									=input_tensors[0].shape[chunk_dim] // chunk_size

								# chunk input tensor into tuples
								lowerCamelCase									=tuple(input_tensor.chunk(_UpperCAmelCase      ,   dim=_UpperCAmelCase					) for input_tensor in input_tensors					)
								# apply forward fn to every tuple
								lowerCamelCase									=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   _lowercase       (    _UpperCAmelCase      ,   _UpperCAmelCase      ,   _UpperCAmelCase      ,   _UpperCAmelCase					)    ->	Tuple[Set[int], torch.LongTensor]:
				lowerCamelCase									=torch.ones(_UpperCAmelCase      ,   _UpperCAmelCase					)
				lowerCamelCase									=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
								lowerCamelCase									=head - sum(1 if h < head else 0 for h in already_pruned_heads					)
								lowerCamelCase									=0
				lowerCamelCase									=mask.view(-1					).contiguous().eq(1					)
				lowerCamelCase									=torch.arange(len(_UpperCAmelCase					)					)[mask].long()
				return heads, index





def   _lowercase       (    *_UpperCAmelCase      ,   _UpperCAmelCase = None					)    ->	Tuple[torch.Tensor, ...]:
				return torch.meshgrid(*_UpperCAmelCase      ,   indexing=_UpperCAmelCase					)





def   _lowercase       (    _UpperCAmelCase					)    ->	Tuple[torch.device, int, int]:
				return tensor.device, storage_ptr(_UpperCAmelCase					), storage_size(_UpperCAmelCase					)












from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


UpperCAmelCase__    :      Optional[Any]       					={
    '''configuration_donut_swin''': ['''DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''DonutSwinConfig'''],
    '''processing_donut''': ['''DonutProcessor'''],
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase__    :      Any       					=[
							    '''DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST''',
							    '''DonutSwinModel''',
							    '''DonutSwinPreTrainedModel''',
							]

try:
							if not is_vision_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase__    :      Tuple       					=['''DonutFeatureExtractor''']
							UpperCAmelCase__    :      List[str]       					=['''DonutImageProcessor''']


if TYPE_CHECKING:
							from .configuration_donut_swin import DONUT_SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP, DonutSwinConfig
							from .processing_donut import DonutProcessor

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_donut_swin import (
														    DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
														    DonutSwinModel,
														    DonutSwinPreTrainedModel,
														)

							try:
														if not is_vision_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .feature_extraction_donut import DonutFeatureExtractor
														from .image_processing_donut import DonutImageProcessor

else:
							import sys

							UpperCAmelCase__    :      int       					=_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





