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


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

_a   :      str       =				logging.get_logger(__name__)

_a   :      Optional[int]       =				version.parse(version.parse(torch.__version__).base_version)

_a   :      str       =				parsed_torch_version_base >= version.parse("""2.0""")
_a   :      Optional[Any]       =				parsed_torch_version_base >= version.parse("""1.12""")
_a   :      Any       =				parsed_torch_version_base >= version.parse("""1.11""")
_a   :      Tuple       =				parsed_torch_version_base < version.parse("""1.11""")
def 					a__ ( a :					List[Any]							,  a :					List[str]							,  a :					Optional[int]							,  a :					Tuple							,  self :					Union[str, Any]		):

  """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 					a__ ( a :					nn.Linear							,  a :					torch.LongTensor							,  a :					int = 0		):

  """simple docstring"""







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






class   _UpperCAmelCase					(       nn.Module):
    def __init__(     self    ,     snake_case_    ,     snake_case_  ):
      super().__init__()
      _snake_case  :    Any            = nf
      _snake_case  :    int            = nn.Parameter(torch.empty(snake_case_    ,     snake_case_  )  )
      _snake_case  :    str            = nn.Parameter(torch.zeros(snake_case_  )  )
      nn.init.normal_(self.weight    ,     std=0.02  )


    def        lowerCamelCase__						(     self    ,     snake_case_  ):
      _snake_case  :    Optional[Any]            = x.size()[:-1] + (self.nf,)
      _snake_case  :    List[Any]            = torch.addmm(self.bias    ,     x.view(-1    ,     x.size(-1  )  )    ,     self.weight  )
      _snake_case  :    Optional[int]            = x.view(snake_case_  )
      return x


def 					a__ ( a :					ConvaD							,  a :					torch.LongTensor							,  a :					int = 1		):

  """simple docstring"""







  _snake_case  :    Optional[int]            = index.to(layer.weight.device		)
  _snake_case  :    List[Any]            = layer.weight.index_select(a							,  a		).clone().detach()
  if dim == 0:
    _snake_case  :    int            = layer.bias.clone().detach()
  else:
    _snake_case  :    Tuple            = layer.bias[index].clone().detach()
  _snake_case  :    int            = list(layer.weight.size()		)
  _snake_case  :    str            = len(a		)
  _snake_case  :    Union[str, Any]            = ConvaD(new_size[1]							,  new_size[0]		).to(layer.weight.device		)
  _snake_case  :    Dict            = False
  new_layer.weight.copy_(W.contiguous()		)
  _snake_case  :    List[Any]            = True
  _snake_case  :    Dict            = False
  new_layer.bias.copy_(b.contiguous()		)
  _snake_case  :    Tuple            = True
  return new_layer
def 					a__ ( a :					Union[nn.Linear, ConvaD]							,  a :					torch.LongTensor							,  a :					Optional[int] = None		):

  """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 					a__ ( a :					Callable[..., torch.Tensor]							,  a :					int							,  a :					int							,  *a :					List[Any]		):

  """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
  _snake_case  :    int            = 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:
    _snake_case  :    Dict            = 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  :    Union[str, Any]            = input_tensors[0].shape[chunk_dim] // chunk_size

    # chunk input tensor into tuples
    _snake_case  :    int            = tuple(input_tensor.chunk(a							,  dim=a		) for input_tensor in input_tensors		)
    # apply forward fn to every tuple
    _snake_case  :    Optional[int]            = 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 					a__ ( a :					List[int]							,  a :					int							,  a :					int							,  a :					Set[int]		):

  """simple docstring"""







  _snake_case  :    str            = torch.ones(a							,  a		)
  _snake_case  :    Tuple            = 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
    _snake_case  :    str            = head - sum(1 if h < head else 0 for h in already_pruned_heads		)
    _snake_case  :    Dict            = 0
  _snake_case  :    List[Any]            = mask.view(-1		).contiguous().eq(1		)
  _snake_case  :    torch.LongTensor            = torch.arange(len(a		)		)[mask].long()
  return heads, index
def 					a__ ( *a :					Union[torch.Tensor, List[torch.Tensor]]							,  a :					Optional[str] = None		):

  """simple docstring"""







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



def 					a__ ( a :					torch.Tensor		):

  """simple docstring"""







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








'''simple docstring'''


from typing import TYPE_CHECKING

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


_lowerCamelCase					:						Dict					     =							{
    '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:
				_lowerCamelCase					:						List[str]					     =							[
				    'DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST',
				    'DonutSwinModel',
				    'DonutSwinPreTrainedModel',
				]

try:
				if not is_vision_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_lowerCamelCase					:						Any					     =							['DonutFeatureExtractor']
				_lowerCamelCase					:						Union[str, Any]					     =							['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

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

