





'''simple docstring'''







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 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      torch.Tensor,      SCREAMING_SNAKE_CASE__   :      torch.Tensor,      SCREAMING_SNAKE_CASE__   :      bool )			->						torch.Tensor:
						# The first operation in a checkpoint can't be in-place, but it's
						# nice to have in-place addition during inference. Thus...
						if not inplace:
												UpperCAmelCase_      :   str      =	ma + ma
						else:
												ma += ma

						return ma





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      torch.Tensor,      SCREAMING_SNAKE_CASE__   :      List[int] )			->						torch.Tensor:
						UpperCAmelCase_      :   Optional[Any]      =	-1 * len(SCREAMING_SNAKE_CASE__ )
						UpperCAmelCase_      :   Optional[Any]      =	list(range(len(tensor.shape[:zero_index] ) ) )
						return tensor.permute(first_inds + [zero_index + i for i in inds] )





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      torch.Tensor,      SCREAMING_SNAKE_CASE__   :      int )			->						torch.Tensor:
						return t.reshape(t.shape[:-no_dims] + (-1,) )





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      torch.Tensor,      SCREAMING_SNAKE_CASE__   :      torch.Tensor,      SCREAMING_SNAKE_CASE__   :      int,      SCREAMING_SNAKE_CASE__   :      float = 1E-4 )			->						torch.Tensor:
						UpperCAmelCase_      :   Union[str, Any]      =	mask.expand(*value.shape )
						return torch.sum(mask * value,      dim=SCREAMING_SNAKE_CASE__ ) / (eps + torch.sum(SCREAMING_SNAKE_CASE__,      dim=SCREAMING_SNAKE_CASE__ ))





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      torch.Tensor,      SCREAMING_SNAKE_CASE__   :      torch.types.Number = 2.31_25,      SCREAMING_SNAKE_CASE__   :      torch.types.Number = 21.68_75,      SCREAMING_SNAKE_CASE__   :      int = 64 )			->						torch.Tensor:
						UpperCAmelCase_      :   Optional[int]      =	torch.linspace(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__,      no_bins - 1,      device=pts.device )
						UpperCAmelCase_      :   Union[str, Any]      =	torch.sqrt(torch.sum((pts.unsqueeze(-2 ) - pts.unsqueeze(-3 )) ** 2,      dim=-1 ) )
						return torch.bucketize(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ )





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Callable[[list], Any],      SCREAMING_SNAKE_CASE__   :      List[dict] )			->						dict:
						UpperCAmelCase_      :   List[Any]      =	dicts[0]
						UpperCAmelCase_      :   Optional[Any]      =	{}
						for k, v in first.items():
												UpperCAmelCase_      :   List[Any]      =	[d[k] for d in dicts]
												if isinstance(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ ):
																		UpperCAmelCase_      :   List[Any]      =	dict_multimap(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ )
												else:
																		UpperCAmelCase_      :   int      =	fn(SCREAMING_SNAKE_CASE__ )

						return new_dict





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





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      torch.Tensor,      SCREAMING_SNAKE_CASE__   :      torch.Tensor,      SCREAMING_SNAKE_CASE__   :      int = 0,      SCREAMING_SNAKE_CASE__   :      int = 0 )			->						torch.Tensor:
						UpperCAmelCase_      :   List[Union[slice, torch.Tensor]]      =	[]
						for i, s in enumerate(data.shape[:no_batch_dims] ):
												UpperCAmelCase_      :   Union[str, Any]      =	torch.arange(SCREAMING_SNAKE_CASE__ )
												UpperCAmelCase_      :   Union[str, Any]      =	r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape ) - i - 1))) )
												ranges.append(SCREAMING_SNAKE_CASE__ )

						UpperCAmelCase_      :   List[Union[slice, torch.Tensor]]      =	[slice(SCREAMING_SNAKE_CASE__ ) for _ in range(len(data.shape ) - no_batch_dims )]
						UpperCAmelCase_      :   Optional[Any]      =	inds
						ranges.extend(SCREAMING_SNAKE_CASE__ )
						# 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(SCREAMING_SNAKE_CASE__ )]


snake_case_  :							List[str]            =	TypeVar("T")





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Callable[[T], Any],      SCREAMING_SNAKE_CASE__   :      Dict[Any, Union[dict, list, tuple, T]],      SCREAMING_SNAKE_CASE__   :      Type[T] )			->						Dict[Any, Union[dict, list, tuple, Any]]:
						UpperCAmelCase_      :   Dict[Any, Union[dict, list, tuple, Any]]      =	{}
						for k, v in dic.items():
												if isinstance(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ ):
																		UpperCAmelCase_      :   Union[str, Any]      =	dict_map(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ )
												else:
																		UpperCAmelCase_      :   int      =	tree_map(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ )

						return new_dict





@overload
def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Callable[[T], Any],      SCREAMING_SNAKE_CASE__   :      T,      SCREAMING_SNAKE_CASE__   :      Type[T] )			->						Any:
						...





@overload
def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Callable[[T], Any],      SCREAMING_SNAKE_CASE__   :      dict,      SCREAMING_SNAKE_CASE__   :      Type[T] )			->						dict:
						...





@overload
def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Callable[[T], Any],      SCREAMING_SNAKE_CASE__   :      list,      SCREAMING_SNAKE_CASE__   :      Type[T] )			->						list:
						...





@overload
def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Callable[[T], Any],      SCREAMING_SNAKE_CASE__   :      tuple,      SCREAMING_SNAKE_CASE__   :      Type[T] )			->						tuple:
						...





def 							lowerCamelCase_   (	SCREAMING_SNAKE_CASE__   :      Dict,      SCREAMING_SNAKE_CASE__   :      Union[str, Any],      SCREAMING_SNAKE_CASE__   :      str )			->						Union[str, Any]:
						if isinstance(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ ):
												return dict_map(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ )
						elif isinstance(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ ):
												return [tree_map(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ ) for x in tree]
						elif isinstance(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ ):
												return tuple(tree_map(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ ) for x in tree )
						elif isinstance(SCREAMING_SNAKE_CASE__,      SCREAMING_SNAKE_CASE__ ):
												return fn(SCREAMING_SNAKE_CASE__ )
						else:
												print(type(SCREAMING_SNAKE_CASE__ ) )
												raise ValueError('''Not supported''' )
snake_case_  :							str            =	partial(tree_map, leaf_type=torch.Tensor)












'''simple docstring'''







from typing import TYPE_CHECKING

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


snake_case_  :							Any            =	{
    "configuration_flava": [
        "FLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "FlavaConfig",
        "FlavaImageCodebookConfig",
        "FlavaImageConfig",
        "FlavaMultimodalConfig",
        "FlavaTextConfig",
    ],
}

try:
					if not is_vision_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					snake_case_  :							List[str]            =	["FlavaFeatureExtractor"]
					snake_case_  :							int            =	["FlavaImageProcessor"]
					snake_case_  :							Tuple            =	["FlavaProcessor"]

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					snake_case_  :							Optional[int]            =	[
					    "FLAVA_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "FlavaForPreTraining",
					    "FlavaImageCodebook",
					    "FlavaImageModel",
					    "FlavaModel",
					    "FlavaMultimodalModel",
					    "FlavaPreTrainedModel",
					    "FlavaTextModel",
					]

if TYPE_CHECKING:
					from .configuration_flava import (
					    FLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP,
					    FlavaConfig,
					    FlavaImageCodebookConfig,
					    FlavaImageConfig,
					    FlavaMultimodalConfig,
					    FlavaTextConfig,
					)

					try:
										if not is_vision_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .feature_extraction_flava import FlavaFeatureExtractor
										from .image_processing_flava import FlavaImageProcessor
										from .processing_flava import FlavaProcessor

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_flava import (
										    FLAVA_PRETRAINED_MODEL_ARCHIVE_LIST,
										    FlavaForPreTraining,
										    FlavaImageCodebook,
										    FlavaImageModel,
										    FlavaModel,
										    FlavaMultimodalModel,
										    FlavaPreTrainedModel,
										    FlavaTextModel,
										)

else:
					import sys

					snake_case_  :							Any            =	_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





