
'''simple docstring'''







import torch
from torch import nn


class 		__UpperCAmelCase   (      nn.Module		):


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





								"""simple docstring"""






								super().__init__()

								_snake_case     		=	n_token
								_snake_case     		=	d_embed
								_snake_case     		=	d_proj

								_snake_case     		=	cutoffs + [n_token]
								_snake_case     		=	[0] + self.cutoffs
								_snake_case     		=	div_val

								_snake_case     		=	self.cutoffs[0]
								_snake_case     		=	len(self.cutoffs						) - 1
								_snake_case     		=	self.shortlist_size + self.n_clusters

								if self.n_clusters > 0:
									_snake_case     		=	nn.Parameter(torch.zeros(self.n_clusters					,     self.d_embed						)						)
									_snake_case     		=	nn.Parameter(torch.zeros(self.n_clusters						)						)

								_snake_case     		=	nn.ModuleList()
								_snake_case     		=	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						)						):
										_snake_case     		=	self.cutoff_ends[i], self.cutoff_ends[i + 1]
										_snake_case     		=	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						)						)

								_snake_case     		=	keep_order


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_						):





								"""simple docstring"""






								if proj is None:
									_snake_case     		=	nn.functional.linear(UpperCAmelCase__					,     UpperCAmelCase__					,     bias=UpperCAmelCase__						)
								else:
									# if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
									_snake_case     		=	nn.functional.linear(UpperCAmelCase__					,     proj.t().contiguous()						)
									_snake_case     		=	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					,     lowerCAmelCase_					,     lowerCAmelCase_=None					,     lowerCAmelCase_=False						):





								"""simple docstring"""






								if labels is not None:
									# Shift so that tokens < n predict n
									_snake_case     		=	hidden[..., :-1, :].contiguous()
									_snake_case     		=	labels[..., 1:].contiguous()
									_snake_case     		=	hidden.view(-1					,     hidden.size(-1						)						)
									_snake_case     		=	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:
									_snake_case     		=	hidden.view(-1					,     hidden.size(-1						)						)

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

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

										weights.append(UpperCAmelCase__						)
										biases.append(UpperCAmelCase__						)

									_snake_case     		=	weights[0], biases[0], self.out_projs[0]

									_snake_case     		=	self._compute_logit(UpperCAmelCase__					,     UpperCAmelCase__					,     UpperCAmelCase__					,     UpperCAmelCase__						)
									_snake_case     		=	nn.functional.log_softmax(UpperCAmelCase__					,     dim=1						)

									if labels is None:
										_snake_case     		=	hidden.new_empty((head_logit.size(0						), self.n_token)						)
									else:
										_snake_case     		=	torch.zeros_like(UpperCAmelCase__					,     dtype=hidden.dtype					,     device=hidden.device						)

									_snake_case     		=	0
									_snake_case     		=	[0] + self.cutoffs
									for i in range(len(UpperCAmelCase__						) - 1						):
										_snake_case     		=	cutoff_values[i], cutoff_values[i + 1]

										if labels is not None:
											_snake_case     		=	(labels >= l_idx) & (labels < r_idx)
											_snake_case     		=	mask_i.nonzero().squeeze()

											if indices_i.numel() == 0:
												continue

											_snake_case     		=	labels.index_select(0					,     UpperCAmelCase__						) - l_idx
											_snake_case     		=	head_logprob.index_select(0					,     UpperCAmelCase__						)
											_snake_case     		=	hidden.index_select(0					,     UpperCAmelCase__						)
										else:
											_snake_case     		=	hidden

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

											_snake_case     		=	self._compute_logit(UpperCAmelCase__					,     UpperCAmelCase__					,     UpperCAmelCase__					,     UpperCAmelCase__						)
											_snake_case     		=	nn.functional.log_softmax(UpperCAmelCase__					,     dim=1						)
											_snake_case     		=	self.cutoffs[0] + i - 1  # No probability for the head cluster
											if labels is not None:
												_snake_case     		=	head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
												    1					,     target_i[:, None]						).squeeze(1						)
											else:
												_snake_case     		=	head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
												_snake_case     		=	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					,     lowerCAmelCase_						):





								"""simple docstring"""






								if self.n_clusters == 0:
									_snake_case     		=	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
									_snake_case     		=	[], []
									for i in range(len(self.cutoffs						)						):
										if self.div_val == 1:
											_snake_case     		=	self.cutoff_ends[i], self.cutoff_ends[i + 1]
											_snake_case     		=	self.out_layers[0].weight[l_idx:r_idx]
											_snake_case     		=	self.out_layers[0].bias[l_idx:r_idx]
										else:
											_snake_case     		=	self.out_layers[i].weight
											_snake_case     		=	self.out_layers[i].bias

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

										weights.append(UpperCAmelCase__						)
										biases.append(UpperCAmelCase__						)

									_snake_case     		=	weights[0], biases[0], self.out_projs[0]
									_snake_case     		=	self._compute_logit(UpperCAmelCase__					,     UpperCAmelCase__					,     UpperCAmelCase__					,     UpperCAmelCase__						)

									_snake_case     		=	hidden.new_empty((head_logit.size(0						), self.n_token)						)
									_snake_case     		=	nn.functional.log_softmax(UpperCAmelCase__					,     dim=1						)

									_snake_case     		=	[0] + self.cutoffs
									for i in range(len(UpperCAmelCase__						) - 1						):
										_snake_case     		=	cutoff_values[i], cutoff_values[i + 1]

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

											_snake_case     		=	self._compute_logit(UpperCAmelCase__					,     UpperCAmelCase__					,     UpperCAmelCase__					,     UpperCAmelCase__						)
											_snake_case     		=	nn.functional.log_softmax(UpperCAmelCase__					,     dim=1						)

											_snake_case     		=	head_logprob[:, -i] + tail_logprob_i
											_snake_case     		=	logprob_i

									return out



'''simple docstring'''







from typing import TYPE_CHECKING

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


lowercase					:							int   =      {
    "configuration_mask2former": [
        "MASK2FORMER_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "Mask2FormerConfig",
    ],
}

try:
					if not is_vision_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowercase					:							Optional[int]   =      ["Mask2FormerImageProcessor"]

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowercase					:							Optional[Any]   =      [
					    "MASK2FORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "Mask2FormerForUniversalSegmentation",
					    "Mask2FormerModel",
					    "Mask2FormerPreTrainedModel",
					]

if TYPE_CHECKING:
					from .configuration_maskaformer import MASK2FORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, MaskaFormerConfig

					try:
										if not is_vision_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .image_processing_maskaformer import MaskaFormerImageProcessor

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_maskaformer import (
										    MASK2FORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
										    MaskaFormerForUniversalSegmentation,
										    MaskaFormerModel,
										    MaskaFormerPreTrainedModel,
										)


else:
					import sys

					lowercase					:							int   =      _LazyModule(__name__, globals()["__file__"], _import_structure)

