
from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types
def     UpperCAmelCase__							(       __magic_name__						:     torch.Tensor							,  __magic_name__						:     torch.Tensor							,  __magic_name__						:     bool  ):


					'''simple docstring'''

					if not inplace:
										lowerCAmelCase		:					Optional[int]  	=     ma + ma
					else:
										ma += ma

					return ma
def     UpperCAmelCase__							(       __magic_name__						:     torch.Tensor							,  __magic_name__						:     List[int]  ):


					'''simple docstring'''

					lowerCAmelCase		:					List[str]  	=     -1 * len(__magic_name__  )
					lowerCAmelCase		:					Tuple  	=     list(range(len(tensor.shape[:zero_index]  )  )  )
					return tensor.permute(first_inds + [zero_index + i for i in inds]  )
def     UpperCAmelCase__							(       __magic_name__						:     torch.Tensor							,  __magic_name__						:     int  ):


					'''simple docstring'''

					return t.reshape(t.shape[:-no_dims] + (-1,)  )
def     UpperCAmelCase__							(       __magic_name__						:     torch.Tensor							,  __magic_name__						:     torch.Tensor							,  __magic_name__						:     int							,  __magic_name__						:     float = 1e-4  ):


					'''simple docstring'''

					lowerCAmelCase		:					Tuple  	=     mask.expand(*value.shape  )
					return torch.sum(mask * value							,  dim=__magic_name__  ) / (eps + torch.sum(__magic_name__							,  dim=__magic_name__  ))
def     UpperCAmelCase__							(       __magic_name__						:     torch.Tensor							,  __magic_name__						:     torch.types.Number = 2.3125							,  __magic_name__						:     torch.types.Number = 21.6875							,  __magic_name__						:     int = 64  ):


					'''simple docstring'''

					lowerCAmelCase		:					Any  	=     torch.linspace(__magic_name__							,  __magic_name__							,  no_bins - 1							,  device=pts.device  )
					lowerCAmelCase		:					Tuple  	=     torch.sqrt(torch.sum((pts.unsqueeze(-2  ) - pts.unsqueeze(-3  )) ** 2							,  dim=-1  )  )
					return torch.bucketize(__magic_name__							,  __magic_name__  )
def     UpperCAmelCase__							(       __magic_name__						:     Callable[[list], Any]							,  __magic_name__						:     List[dict]  ):


					'''simple docstring'''

					lowerCAmelCase		:					Optional[Any]  	=     dicts[0]
					lowerCAmelCase		:					int  	=     {}
					for k, v in first.items():
										lowerCAmelCase		:					Dict  	=     [d[k] for d in dicts]
										if isinstance(__magic_name__							,  __magic_name__  ):
															lowerCAmelCase		:					Any  	=     dict_multimap(__magic_name__							,  __magic_name__  )
										else:
															lowerCAmelCase		:					List[str]  	=     fn(__magic_name__  )

					return new_dict
def     UpperCAmelCase__							(       __magic_name__						:     torch.Tensor							,  __magic_name__						:     torch.Tensor  ):


					'''simple docstring'''

					lowerCAmelCase		:					List[Any]  	=     v_bins.view(((1,) * len(x.shape  )) + (len(__magic_name__  ),)  )
					lowerCAmelCase		:					Union[str, Any]  	=     x[..., None] - reshaped_bins
					lowerCAmelCase		:					Union[str, Any]  	=     torch.argmin(torch.abs(__magic_name__  )							,  dim=-1  )
					return nn.functional.one_hot(__magic_name__							,  num_classes=len(__magic_name__  )  ).float()
def     UpperCAmelCase__							(       __magic_name__						:     torch.Tensor							,  __magic_name__						:     torch.Tensor							,  __magic_name__						:     int = 0							,  __magic_name__						:     int = 0  ):


					'''simple docstring'''

					lowerCAmelCase		:					List[Union[slice, torch.Tensor]]  	=     []
					for i, s in enumerate(data.shape[:no_batch_dims]  ):
										lowerCAmelCase		:					Any  	=     torch.arange(__magic_name__  )
										lowerCAmelCase		:					int  	=     r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape  ) - i - 1)))  )
										ranges.append(__magic_name__  )

					lowerCAmelCase		:					List[Union[slice, torch.Tensor]]  	=     [slice(__magic_name__  ) for _ in range(len(data.shape  ) - no_batch_dims  )]
					lowerCAmelCase		:					List[Any]  	=     inds
					ranges.extend(__magic_name__  )
					# Matt note: Editing this to get around the behaviour of using a list as an array index changing
					# in recent Numpy versions
					return data[tuple(__magic_name__  )]


__SCREAMING_SNAKE_CASE			:   Optional[int]         =     TypeVar('T')
def     UpperCAmelCase__							(       __magic_name__						:     Callable[[T], Any]							,  __magic_name__						:     Dict[Any, Union[dict, list, tuple, T]]							,  __magic_name__						:     Type[T]  ):


					'''simple docstring'''

					lowerCAmelCase		:					Dict[Any, Union[dict, list, tuple, Any]]  	=     {}
					for k, v in dic.items():
										if isinstance(__magic_name__							,  __magic_name__  ):
															lowerCAmelCase		:					Optional[int]  	=     dict_map(__magic_name__							,  __magic_name__							,  __magic_name__  )
										else:
															lowerCAmelCase		:					Tuple  	=     tree_map(__magic_name__							,  __magic_name__							,  __magic_name__  )

					return new_dict
@overload
def     UpperCAmelCase__							(       __magic_name__						:     Callable[[T], Any]							,  __magic_name__						:     T							,  __magic_name__						:     Type[T]  ):


					'''simple docstring'''

					...
@overload
def     UpperCAmelCase__							(       __magic_name__						:     Callable[[T], Any]							,  __magic_name__						:     dict							,  __magic_name__						:     Type[T]  ):


					'''simple docstring'''

					...
@overload
def     UpperCAmelCase__							(       __magic_name__						:     Callable[[T], Any]							,  __magic_name__						:     list							,  __magic_name__						:     Type[T]  ):


					'''simple docstring'''

					...
@overload
def     UpperCAmelCase__							(       __magic_name__						:     Callable[[T], Any]							,  __magic_name__						:     tuple							,  __magic_name__						:     Type[T]  ):


					'''simple docstring'''

					...
def     UpperCAmelCase__							(       __magic_name__						:     Optional[int]							,  __magic_name__						:     Dict							,  __magic_name__						:     List[Any]  ):


					'''simple docstring'''

					if isinstance(__magic_name__							,  __magic_name__  ):
										return dict_map(__magic_name__							,  __magic_name__							,  __magic_name__  )
					elif isinstance(__magic_name__							,  __magic_name__  ):
										return [tree_map(__magic_name__							,  __magic_name__							,  __magic_name__  ) for x in tree]
					elif isinstance(__magic_name__							,  __magic_name__  ):
										return tuple(tree_map(__magic_name__							,  __magic_name__							,  __magic_name__  ) for x in tree  )
					elif isinstance(__magic_name__							,  __magic_name__  ):
										return fn(__magic_name__  )
					else:
										print(type(__magic_name__  )  )
										raise ValueError('''Not supported'''  )
__SCREAMING_SNAKE_CASE			:   Union[str, Any]         =     partial(tree_map, leaf_type=torch.Tensor)


from typing import TYPE_CHECKING

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


__SCREAMING_SNAKE_CASE			:   List[str]         =     {'configuration_vilt': ['VILT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'ViltConfig']}

try:
		if not is_vision_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__SCREAMING_SNAKE_CASE			:   Any         =     ['ViltFeatureExtractor']
		__SCREAMING_SNAKE_CASE			:   Optional[int]         =     ['ViltImageProcessor']
		__SCREAMING_SNAKE_CASE			:   str         =     ['ViltProcessor']

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__SCREAMING_SNAKE_CASE			:   Dict         =     [
		    'VILT_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'ViltForImageAndTextRetrieval',
		    'ViltForImagesAndTextClassification',
		    'ViltForTokenClassification',
		    'ViltForMaskedLM',
		    'ViltForQuestionAnswering',
		    'ViltLayer',
		    'ViltModel',
		    'ViltPreTrainedModel',
		]


if TYPE_CHECKING:
		from .configuration_vilt import VILT_PRETRAINED_CONFIG_ARCHIVE_MAP, ViltConfig

		try:
				if not is_vision_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .feature_extraction_vilt import ViltFeatureExtractor
				from .image_processing_vilt import ViltImageProcessor
				from .processing_vilt import ViltProcessor

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_vilt import (
				    VILT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    ViltForImageAndTextRetrieval,
				    ViltForImagesAndTextClassification,
				    ViltForMaskedLM,
				    ViltForQuestionAnswering,
				    ViltForTokenClassification,
				    ViltLayer,
				    ViltModel,
				    ViltPreTrainedModel,
				)


else:
		import sys

		__SCREAMING_SNAKE_CASE			:   Optional[Any]         =     _LazyModule(__name__, globals()['__file__'], _import_structure)
