


'''simple docstring'''
import torch
from torch import nn
class    lowercase	(							nn.Module					):

		"""simple docstring"""

		def __init__(		self    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_=1    ,   UpperCamelCase_=False							):

							'''simple docstring'''


							super().__init__()

							UpperCamelCase__   :Union[str, Any]												= n_token
							UpperCamelCase__   :Union[str, Any]												= d_embed
							UpperCamelCase__   :List[Any]												= d_proj

							UpperCamelCase__   :Optional[Any]												= cutoffs + [n_token]
							UpperCamelCase__   :Any												= [0] + self.cutoffs
							UpperCamelCase__   :Any												= div_val

							UpperCamelCase__   :Optional[int]												= self.cutoffs[0]
							UpperCamelCase__   :List[Any]												= len(self.cutoffs							) - 1
							UpperCamelCase__   :Dict												= self.shortlist_size + self.n_clusters

							if self.n_clusters > 0:
												UpperCamelCase__   :int												= nn.Parameter(torch.zeros(self.n_clusters    ,   self.d_embed							)							)
												UpperCamelCase__   :Union[str, Any]												= nn.Parameter(torch.zeros(self.n_clusters							)							)

							UpperCamelCase__   :Any												= nn.ModuleList()
							UpperCamelCase__   :Optional[int]												= nn.ParameterList()

							if div_val == 1:
												for i in range(len(self.cutoffs							)							):
																	if d_proj != d_embed:
																						self.out_projs.append(nn.Parameter(torch.FloatTensor(UpperCamelCase_    ,   UpperCamelCase_							)							)							)
																	else:
																						self.out_projs.append(UpperCamelCase_							)

												self.out_layers.append(nn.Linear(UpperCamelCase_    ,   UpperCamelCase_							)							)
							else:
												for i in range(len(self.cutoffs							)							):
																	UpperCamelCase__					,				UpperCamelCase__   :str												= self.cutoff_ends[i], self.cutoff_ends[i + 1]
																	UpperCamelCase__   :Union[str, Any]												= d_embed // (div_val**i)

																	self.out_projs.append(nn.Parameter(torch.FloatTensor(UpperCamelCase_    ,   UpperCamelCase_							)							)							)

																	self.out_layers.append(nn.Linear(UpperCamelCase_    ,   r_idx - l_idx							)							)

							UpperCamelCase__   :Optional[int]												= keep_order




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_							):

							'''simple docstring'''


							if proj is None:
												UpperCamelCase__   :int												= nn.functional.linear(UpperCamelCase_    ,   UpperCamelCase_    ,   bias=UpperCamelCase_							)
							else:
												# if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
												UpperCamelCase__   :int												= nn.functional.linear(UpperCamelCase_    ,   proj.t().contiguous()							)
												UpperCamelCase__   :Any												= nn.functional.linear(UpperCamelCase_    ,   UpperCamelCase_    ,   bias=UpperCamelCase_							)
												# else:
												#     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
												#     if bias is not None:
												#         logit = logit + bias

							return logit




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_    ,   UpperCamelCase_=None    ,   UpperCamelCase_=False							):

							'''simple docstring'''


							if labels is not None:
												# Shift so that tokens < n predict n
												UpperCamelCase__   :int												= hidden[..., :-1, :].contiguous()
												UpperCamelCase__   :int												= labels[..., 1:].contiguous()
												UpperCamelCase__   :Dict												= hidden.view(-1    ,   hidden.size(-1							)							)
												UpperCamelCase__   :List[Any]												= labels.view(-1							)
												if hidden.size(0							) != labels.size(0							):
																	raise RuntimeError('''Input and labels should have the same size in the batch dimension.'''							)
							else:
												UpperCamelCase__   :List[str]												= hidden.view(-1    ,   hidden.size(-1							)							)

							if self.n_clusters == 0:
												UpperCamelCase__   :List[Any]												= self._compute_logit(UpperCamelCase_    ,   self.out_layers[0].weight    ,   self.out_layers[0].bias    ,   self.out_projs[0]							)
												if labels is not None:
																	UpperCamelCase__   :Tuple												= labels != -100
																	UpperCamelCase__   :List[str]												= torch.zeros_like(UpperCamelCase_    ,   dtype=hidden.dtype    ,   device=hidden.device							)
																	UpperCamelCase__   :List[str]												= (
																	    -nn.functional.log_softmax(UpperCamelCase_    ,   dim=-1							)[mask].gather(1    ,   labels[mask].unsqueeze(1							)							).squeeze(1							)
																	)
												else:
																	UpperCamelCase__   :Dict												= nn.functional.log_softmax(UpperCamelCase_    ,   dim=-1							)
							else:
												# construct weights and biases
												UpperCamelCase__					,				UpperCamelCase__   :int												= [], []
												for i in range(len(self.cutoffs							)							):
																	if self.div_val == 1:
																						UpperCamelCase__					,				UpperCamelCase__   :Any												= self.cutoff_ends[i], self.cutoff_ends[i + 1]
																						UpperCamelCase__   :Union[str, Any]												= self.out_layers[0].weight[l_idx:r_idx]
																						UpperCamelCase__   :Tuple												= self.out_layers[0].bias[l_idx:r_idx]
																	else:
																						UpperCamelCase__   :List[Any]												= self.out_layers[i].weight
																						UpperCamelCase__   :Tuple												= self.out_layers[i].bias

																	if i == 0:
																						UpperCamelCase__   :List[Any]												= torch.cat([weight_i, self.cluster_weight]    ,   dim=0							)
																						UpperCamelCase__   :List[Any]												= torch.cat([bias_i, self.cluster_bias]    ,   dim=0							)

																	weights.append(UpperCamelCase_							)
																	biases.append(UpperCamelCase_							)

												UpperCamelCase__					,				UpperCamelCase__					,				UpperCamelCase__   :List[Any]												= weights[0], biases[0], self.out_projs[0]

												UpperCamelCase__   :List[Any]												= self._compute_logit(UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_							)
												UpperCamelCase__   :Union[str, Any]												= nn.functional.log_softmax(UpperCamelCase_    ,   dim=1							)

												if labels is None:
																	UpperCamelCase__   :Dict												= hidden.new_empty((head_logit.size(0							), self.n_token)							)
												else:
																	UpperCamelCase__   :Optional[Any]												= torch.zeros_like(UpperCamelCase_    ,   dtype=hidden.dtype    ,   device=hidden.device							)

												UpperCamelCase__   :Any												= 0
												UpperCamelCase__   :List[str]												= [0] + self.cutoffs
												for i in range(len(UpperCamelCase_							) - 1							):
																	UpperCamelCase__					,				UpperCamelCase__   :Any												= cutoff_values[i], cutoff_values[i + 1]

																	if labels is not None:
																						UpperCamelCase__   :List[Any]												= (labels >= l_idx) & (labels < r_idx)
																						UpperCamelCase__   :List[Any]												= mask_i.nonzero().squeeze()

																						if indices_i.numel() == 0:
																											continue

																						UpperCamelCase__   :Optional[Any]												= labels.index_select(0    ,   UpperCamelCase_							) - l_idx
																						UpperCamelCase__   :List[str]												= head_logprob.index_select(0    ,   UpperCamelCase_							)
																						UpperCamelCase__   :Tuple												= hidden.index_select(0    ,   UpperCamelCase_							)
																	else:
																						UpperCamelCase__   :Any												= hidden

																	if i == 0:
																						if labels is not None:
																											UpperCamelCase__   :str												= head_logprob_i.gather(1    ,   target_i[:, None]							).squeeze(1							)
																						else:
																											UpperCamelCase__   :List[Any]												= head_logprob[:, : self.cutoffs[0]]
																	else:
																						UpperCamelCase__					,				UpperCamelCase__					,				UpperCamelCase__   :Optional[Any]												= weights[i], biases[i], self.out_projs[i]

																						UpperCamelCase__   :Optional[Any]												= self._compute_logit(UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_							)
																						UpperCamelCase__   :Any												= nn.functional.log_softmax(UpperCamelCase_    ,   dim=1							)
																						UpperCamelCase__   :int												= self.cutoffs[0] + i - 1  # No probability for the head cluster
																						if labels is not None:
																											UpperCamelCase__   :List[str]												= head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
																											    1    ,   target_i[:, None]							).squeeze(1							)
																						else:
																											UpperCamelCase__   :Any												= head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
																											UpperCamelCase__   :List[str]												= logprob_i

																	if labels is not None:
																						if (hasattr(self    ,   '''keep_order'''							) and self.keep_order) or keep_order:
																											out.index_copy_(0    ,   UpperCamelCase_    ,   -logprob_i							)
																						else:
																											out[offset : offset + logprob_i.size(0							)].copy_(-logprob_i							)
																						offset += logprob_i.size(0							)

							return out




		def     lowerCAmelCase__				(		self    ,   UpperCamelCase_							):

							'''simple docstring'''


							if self.n_clusters == 0:
												UpperCamelCase__   :int												= self._compute_logit(UpperCamelCase_    ,   self.out_layers[0].weight    ,   self.out_layers[0].bias    ,   self.out_projs[0]							)
												return nn.functional.log_softmax(UpperCamelCase_    ,   dim=-1							)
							else:
												# construct weights and biases
												UpperCamelCase__					,				UpperCamelCase__   :List[str]												= [], []
												for i in range(len(self.cutoffs							)							):
																	if self.div_val == 1:
																						UpperCamelCase__					,				UpperCamelCase__   :int												= self.cutoff_ends[i], self.cutoff_ends[i + 1]
																						UpperCamelCase__   :Union[str, Any]												= self.out_layers[0].weight[l_idx:r_idx]
																						UpperCamelCase__   :Any												= self.out_layers[0].bias[l_idx:r_idx]
																	else:
																						UpperCamelCase__   :int												= self.out_layers[i].weight
																						UpperCamelCase__   :Tuple												= self.out_layers[i].bias

																	if i == 0:
																						UpperCamelCase__   :Optional[int]												= torch.cat([weight_i, self.cluster_weight]    ,   dim=0							)
																						UpperCamelCase__   :str												= torch.cat([bias_i, self.cluster_bias]    ,   dim=0							)

																	weights.append(UpperCamelCase_							)
																	biases.append(UpperCamelCase_							)

												UpperCamelCase__					,				UpperCamelCase__					,				UpperCamelCase__   :Union[str, Any]												= weights[0], biases[0], self.out_projs[0]
												UpperCamelCase__   :int												= self._compute_logit(UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_							)

												UpperCamelCase__   :str												= hidden.new_empty((head_logit.size(0							), self.n_token)							)
												UpperCamelCase__   :Any												= nn.functional.log_softmax(UpperCamelCase_    ,   dim=1							)

												UpperCamelCase__   :Optional[Any]												= [0] + self.cutoffs
												for i in range(len(UpperCamelCase_							) - 1							):
																	UpperCamelCase__					,				UpperCamelCase__   :Union[str, Any]												= cutoff_values[i], cutoff_values[i + 1]

																	if i == 0:
																						UpperCamelCase__   :str												= head_logprob[:, : self.cutoffs[0]]
																	else:
																						UpperCamelCase__					,				UpperCamelCase__					,				UpperCamelCase__   :int												= weights[i], biases[i], self.out_projs[i]

																						UpperCamelCase__   :Dict												= self._compute_logit(UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_							)
																						UpperCamelCase__   :Tuple												= nn.functional.log_softmax(UpperCamelCase_    ,   dim=1							)

																						UpperCamelCase__   :str												= head_logprob[:, -i] + tail_logprob_i
																						UpperCamelCase__   :str												= logprob_i

												return out



'''simple docstring'''
from typing import TYPE_CHECKING

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


__snake_case							      =     {'''configuration_dpt''': ['''DPT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''DPTConfig''']}

try:
		if not is_vision_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__snake_case							      =     ['''DPTFeatureExtractor''']
		__snake_case							      =     ['''DPTImageProcessor''']

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__snake_case							      =     [
		    '''DPT_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''DPTForDepthEstimation''',
		    '''DPTForSemanticSegmentation''',
		    '''DPTModel''',
		    '''DPTPreTrainedModel''',
		]


if TYPE_CHECKING:
		from .configuration_dpt import DPT_PRETRAINED_CONFIG_ARCHIVE_MAP, DPTConfig

		try:
				if not is_vision_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .feature_extraction_dpt import DPTFeatureExtractor
				from .image_processing_dpt import DPTImageProcessor

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_dpt import (
				    DPT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    DPTForDepthEstimation,
				    DPTForSemanticSegmentation,
				    DPTModel,
				    DPTPreTrainedModel,
				)


else:
		import sys

		__snake_case							      =     _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)