





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     =			[nn.LayerNorm]

__lowerCAmelCase     =			logging.get_logger(__name__)

__lowerCAmelCase     =			version.parse(version.parse(torch.__version__).base_version)

__lowerCAmelCase     =			parsed_torch_version_base >= version.parse("2.0")
__lowerCAmelCase     =			parsed_torch_version_base >= version.parse("1.12")
__lowerCAmelCase     =			parsed_torch_version_base >= version.parse("1.11")
__lowerCAmelCase     =			parsed_torch_version_base < version.parse("1.11")

def     a						(      _lowercase   :     Any ,						_lowercase   :     int ,						_lowercase   :     str ,						_lowercase   :     int ,						self   :     Tuple       )     ->int:
		from torch import _softmax_backward_data

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

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

class   _UpperCAmelCase	(   nn.Module					):




					def __init__(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase       )						->  str:
							super().__init__()
							UpperCamelCase__         =	nf
							UpperCamelCase__         =	nn.Parameter(torch.empty(__lowerCAmelCase    ,      __lowerCAmelCase       )       )
							UpperCamelCase__         =	nn.Parameter(torch.zeros(__lowerCAmelCase       )       )
							nn.init.normal_(self.weight    ,      std=0.0_2       )




					def 	snake_case_		(						self    ,      __lowerCAmelCase       )						->  int:
							UpperCamelCase__         =	x.size()[:-1] + (self.nf,)
							UpperCamelCase__         =	torch.addmm(self.bias    ,      x.view(-1    ,      x.size(-1       )       )    ,      self.weight       )
							UpperCamelCase__         =	x.view(__lowerCAmelCase       )
							return x



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

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

def     a						(      _lowercase   :     Callable[..., torch.Tensor] ,						_lowercase   :     int ,						_lowercase   :     int ,						*_lowercase   :     Optional[int]       )     ->torch.Tensor:
		assert len(_lowercase       ) > 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
		UpperCamelCase__         =	len(inspect.signature(_lowercase       ).parameters       )
		if num_args_in_forward_chunk_fn != len(_lowercase       ):
				raise ValueError(
				    F"""forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(_lowercase       )} input """
				    """tensors are given"""       )

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

				UpperCamelCase__         =	input_tensors[0].shape[chunk_dim] // chunk_size

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

		return forward_fn(*_lowercase       )

def     a						(      _lowercase   :     List[int] ,						_lowercase   :     int ,						_lowercase   :     int ,						_lowercase   :     Set[int]       )     ->Tuple[Set[int], torch.LongTensor]:
		UpperCamelCase__         =	torch.ones(_lowercase ,						_lowercase       )
		UpperCamelCase__         =	set(_lowercase       ) - 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
				UpperCamelCase__         =	head - sum(1 if h < head else 0 for h in already_pruned_heads       )
				UpperCamelCase__         =	0
		UpperCamelCase__         =	mask.view(-1       ).contiguous().eq(1       )
		UpperCamelCase__         =	torch.arange(len(_lowercase       )       )[mask].long()
		return heads, index

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

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












from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    get_resize_output_image_size,
    normalize,
    pad,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging
from ...utils.import_utils import is_vision_available


__lowerCAmelCase     =			logging.get_logger(__name__)


if is_vision_available():
				import PIL




class   _UpperCAmelCase	(   _lowercase					):
					__A          =   ["pixel_values"]




					def __init__(						self    ,      __lowerCAmelCase = True    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = PILImageResampling.BILINEAR    ,      __lowerCAmelCase = True    ,      __lowerCAmelCase = False    ,      __lowerCAmelCase = True    ,      __lowerCAmelCase = True    ,      __lowerCAmelCase = 1 / 255    ,      __lowerCAmelCase = True    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = None    ,      **__lowerCAmelCase    ,      )						->  None:
							super().__init__(**__lowerCAmelCase       )

							UpperCamelCase__         =	size if size is not None else {"""height""": 2_560, """width""": 1_920}
							if isinstance(__lowerCAmelCase    ,      (tuple, list)       ):
									# The previous feature extractor size parameter was in (width, height) format
									UpperCamelCase__         =	size[::-1]
							UpperCamelCase__         =	get_size_dict(__lowerCAmelCase       )

							UpperCamelCase__         =	do_resize
							UpperCamelCase__         =	size
							UpperCamelCase__         =	resample
							UpperCamelCase__         =	do_thumbnail
							UpperCamelCase__         =	do_align_long_axis
							UpperCamelCase__         =	do_pad
							UpperCamelCase__         =	do_rescale
							UpperCamelCase__         =	rescale_factor
							UpperCamelCase__         =	do_normalize
							UpperCamelCase__         =	image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
							UpperCamelCase__         =	image_std if image_std is not None else IMAGENET_STANDARD_STD




					def 	snake_case_		(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase = None       )						->  np.ndarray:
							UpperCamelCase__  ,			UpperCamelCase__         =	get_image_size(__lowerCAmelCase       )
							UpperCamelCase__  ,			UpperCamelCase__         =	size["""height"""], size["""width"""]

							if (output_width < output_height and input_width > input_height) or (
							    output_width > output_height and input_width < input_height
							):
									UpperCamelCase__         =	np.rotaa(__lowerCAmelCase    ,      3       )

							if data_format is not None:
									UpperCamelCase__         =	to_channel_dimension_format(__lowerCAmelCase    ,      __lowerCAmelCase       )

							return image




					def 	snake_case_		(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase = False    ,      __lowerCAmelCase = None    ,      )						->  np.ndarray:
							UpperCamelCase__  ,			UpperCamelCase__         =	size["""height"""], size["""width"""]
							UpperCamelCase__  ,			UpperCamelCase__         =	get_image_size(__lowerCAmelCase       )

							UpperCamelCase__         =	output_width - input_width
							UpperCamelCase__         =	output_height - input_height

							if random_padding:
									UpperCamelCase__         =	np.random.randint(low=0    ,      high=delta_height + 1       )
									UpperCamelCase__         =	np.random.randint(low=0    ,      high=delta_width + 1       )
							else:
									UpperCamelCase__         =	delta_height // 2
									UpperCamelCase__         =	delta_width // 2

							UpperCamelCase__         =	delta_height - pad_top
							UpperCamelCase__         =	delta_width - pad_left

							UpperCamelCase__         =	((pad_top, pad_bottom), (pad_left, pad_right))
							return pad(__lowerCAmelCase    ,      __lowerCAmelCase    ,      data_format=__lowerCAmelCase       )




					def 	snake_case_		(						self    ,      *__lowerCAmelCase    ,      **__lowerCAmelCase       )						->  List[Any]:
							logger.info("""pad is deprecated and will be removed in version 4.27. Please use pad_image instead."""       )
							return self.pad_image(*__lowerCAmelCase    ,      **__lowerCAmelCase       )




					def 	snake_case_		(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase = PILImageResampling.BICUBIC    ,      __lowerCAmelCase = None    ,      **__lowerCAmelCase    ,      )						->  np.ndarray:
							UpperCamelCase__  ,			UpperCamelCase__         =	get_image_size(__lowerCAmelCase       )
							UpperCamelCase__  ,			UpperCamelCase__         =	size["""height"""], size["""width"""]

							# We always resize to the smallest of either the input or output size.
							UpperCamelCase__         =	min(__lowerCAmelCase    ,      __lowerCAmelCase       )
							UpperCamelCase__         =	min(__lowerCAmelCase    ,      __lowerCAmelCase       )

							if height == input_height and width == input_width:
									return image

							if input_height > input_width:
									UpperCamelCase__         =	int(input_width * height / input_height       )
							elif input_width > input_height:
									UpperCamelCase__         =	int(input_height * width / input_width       )

							return resize(
							    __lowerCAmelCase    ,      size=(height, width)    ,      resample=__lowerCAmelCase    ,      reducing_gap=2.0    ,      data_format=__lowerCAmelCase    ,      **__lowerCAmelCase       )




					def 	snake_case_		(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase = PILImageResampling.BICUBIC    ,      __lowerCAmelCase = None    ,      **__lowerCAmelCase    ,      )						->  np.ndarray:
							UpperCamelCase__         =	get_size_dict(__lowerCAmelCase       )
							UpperCamelCase__         =	min(size["""height"""]    ,      size["""width"""]       )
							UpperCamelCase__         =	get_resize_output_image_size(__lowerCAmelCase    ,      size=__lowerCAmelCase    ,      default_to_square=__lowerCAmelCase       )
							UpperCamelCase__         =	resize(__lowerCAmelCase    ,      size=__lowerCAmelCase    ,      resample=__lowerCAmelCase    ,      data_format=__lowerCAmelCase    ,      **__lowerCAmelCase       )
							return resized_image




					def 	snake_case_		(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase = None    ,      **__lowerCAmelCase    ,      )						->  List[str]:
							return rescale(__lowerCAmelCase    ,      scale=__lowerCAmelCase    ,      data_format=__lowerCAmelCase    ,      **__lowerCAmelCase       )




					def 	snake_case_		(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase    ,      __lowerCAmelCase = None    ,      **__lowerCAmelCase    ,      )						->  np.ndarray:
							return normalize(__lowerCAmelCase    ,      mean=__lowerCAmelCase    ,      std=__lowerCAmelCase    ,      data_format=__lowerCAmelCase    ,      **__lowerCAmelCase       )




					def 	snake_case_		(						self    ,      __lowerCAmelCase    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = False    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = None    ,      __lowerCAmelCase = ChannelDimension.FIRST    ,      **__lowerCAmelCase    ,      )						->  PIL.Image.Image:
							UpperCamelCase__         =	do_resize if do_resize is not None else self.do_resize
							UpperCamelCase__         =	size if size is not None else self.size
							if isinstance(__lowerCAmelCase    ,      (tuple, list)       ):
									# Previous feature extractor had size in (width, height) format
									UpperCamelCase__         =	size[::-1]
							UpperCamelCase__         =	get_size_dict(__lowerCAmelCase       )
							UpperCamelCase__         =	resample if resample is not None else self.resample
							UpperCamelCase__         =	do_thumbnail if do_thumbnail is not None else self.do_thumbnail
							UpperCamelCase__         =	do_align_long_axis if do_align_long_axis is not None else self.do_align_long_axis
							UpperCamelCase__         =	do_pad if do_pad is not None else self.do_pad
							UpperCamelCase__         =	do_rescale if do_rescale is not None else self.do_rescale
							UpperCamelCase__         =	rescale_factor if rescale_factor is not None else self.rescale_factor
							UpperCamelCase__         =	do_normalize if do_normalize is not None else self.do_normalize
							UpperCamelCase__         =	image_mean if image_mean is not None else self.image_mean
							UpperCamelCase__         =	image_std if image_std is not None else self.image_std

							UpperCamelCase__         =	make_list_of_images(__lowerCAmelCase       )

							if not valid_images(__lowerCAmelCase       ):
									raise ValueError(
									    """Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, """
									    """torch.Tensor, tf.Tensor or jax.ndarray."""       )

							if do_resize and size is None:
									raise ValueError("""Size must be specified if do_resize is True."""       )

							if do_rescale and rescale_factor is None:
									raise ValueError("""Rescale factor must be specified if do_rescale is True."""       )

							if do_pad and size is None:
									raise ValueError("""Size must be specified if do_pad is True."""       )

							if do_normalize and (image_mean is None or image_std is None):
									raise ValueError("""Image mean and std must be specified if do_normalize is True."""       )

							# All transformations expect numpy arrays.
							UpperCamelCase__         =	[to_numpy_array(__lowerCAmelCase       ) for image in images]

							if do_align_long_axis:
									UpperCamelCase__         =	[self.align_long_axis(__lowerCAmelCase    ,      size=__lowerCAmelCase       ) for image in images]

							if do_resize:
									UpperCamelCase__         =	[self.resize(image=__lowerCAmelCase    ,      size=__lowerCAmelCase    ,      resample=__lowerCAmelCase       ) for image in images]

							if do_thumbnail:
									UpperCamelCase__         =	[self.thumbnail(image=__lowerCAmelCase    ,      size=__lowerCAmelCase       ) for image in images]

							if do_pad:
									UpperCamelCase__         =	[self.pad_image(image=__lowerCAmelCase    ,      size=__lowerCAmelCase    ,      random_padding=__lowerCAmelCase       ) for image in images]

							if do_rescale:
									UpperCamelCase__         =	[self.rescale(image=__lowerCAmelCase    ,      scale=__lowerCAmelCase       ) for image in images]

							if do_normalize:
									UpperCamelCase__         =	[self.normalize(image=__lowerCAmelCase    ,      mean=__lowerCAmelCase    ,      std=__lowerCAmelCase       ) for image in images]

							UpperCamelCase__         =	[to_channel_dimension_format(__lowerCAmelCase    ,      __lowerCAmelCase       ) for image in images]

							UpperCamelCase__         =	{"""pixel_values""": images}
							return BatchFeature(data=__lowerCAmelCase    ,      tensor_type=__lowerCAmelCase       )





