





'''simple docstring'''

import torch
from torch import nn
class     __lowerCAmelCase					(   nn.Module			):






       """simple docstring"""





       def __init__(							self       :		Tuple	,			lowerCAmelCase__       :		Dict	,			lowerCAmelCase__       :		List[Any]	,			lowerCAmelCase__       :		List[Any]	,			lowerCAmelCase__       :		Tuple	,			lowerCAmelCase__       :		Union[str, Any]=1	,			lowerCAmelCase__       :		Any=False					)   ->					Tuple:

            '''simple docstring'''




            super().__init__()

            _UpperCamelCase       =	n_token
            _UpperCamelCase       =	d_embed
            _UpperCamelCase       =	d_proj

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

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

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

            _UpperCamelCase       =	nn.ModuleList()
            _UpperCamelCase       =	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(lowerCAmelCase__	,			lowerCAmelCase__					)					)					)
                      else:
                           self.out_projs.append(lowerCAmelCase__					)

                 self.out_layers.append(nn.Linear(lowerCAmelCase__	,			lowerCAmelCase__					)					)
            else:
                 for i in range(len(self.cutoffs					)					):
                      _UpperCamelCase       =	self.cutoff_ends[i], self.cutoff_ends[i + 1]
                      _UpperCamelCase       =	d_embed // (div_val**i)

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

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

            _UpperCamelCase       =	keep_order





       def snake_case__						(							self       :		Tuple	,			lowerCAmelCase__       :		List[Any]	,			lowerCAmelCase__       :		List[Any]	,			lowerCAmelCase__       :		Union[str, Any]	,			lowerCAmelCase__       :		Union[str, Any]					)   ->					Any:

            '''simple docstring'''




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

            return logit





       def snake_case__						(							self       :		str	,			lowerCAmelCase__       :		Union[str, Any]	,			lowerCAmelCase__       :		Union[str, Any]=None	,			lowerCAmelCase__       :		Dict=False					)   ->					Optional[int]:

            '''simple docstring'''




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

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

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

                      weights.append(lowerCAmelCase__					)
                      biases.append(lowerCAmelCase__					)

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

                 _UpperCamelCase       =	self._compute_logit(lowerCAmelCase__	,			lowerCAmelCase__	,			lowerCAmelCase__	,			lowerCAmelCase__					)
                 _UpperCamelCase       =	nn.functional.log_softmax(lowerCAmelCase__	,			dim=1					)

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

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

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

                           if indices_i.numel() == 0:
                                continue

                           _UpperCamelCase       =	labels.index_select(0	,			lowerCAmelCase__					) - l_idx
                           _UpperCamelCase       =	head_logprob.index_select(0	,			lowerCAmelCase__					)
                           _UpperCamelCase       =	hidden.index_select(0	,			lowerCAmelCase__					)
                      else:
                           _UpperCamelCase       =	hidden

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

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

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

            return out







       def snake_case__						(							self       :		Any	,			lowerCAmelCase__       :		Tuple					)   ->					Union[str, Any]:

            '''simple docstring'''




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

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

                      weights.append(lowerCAmelCase__					)
                      biases.append(lowerCAmelCase__					)

                 _UpperCamelCase       =	weights[0], biases[0], self.out_projs[0]
                 _UpperCamelCase       =	self._compute_logit(lowerCAmelCase__	,			lowerCAmelCase__	,			lowerCAmelCase__	,			lowerCAmelCase__					)

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

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

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

                           _UpperCamelCase       =	self._compute_logit(lowerCAmelCase__	,			lowerCAmelCase__	,			lowerCAmelCase__	,			lowerCAmelCase__					)
                           _UpperCamelCase       =	nn.functional.log_softmax(lowerCAmelCase__	,			dim=1					)

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

                 return out










'''simple docstring'''

from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowercase__						:      Optional[Any]	=    {
    'configuration_autoformer': [
        'AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'AutoformerConfig',
    ],
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowercase__						:      List[Any]	=    [
       'AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
       'AutoformerForPrediction',
       'AutoformerModel',
       'AutoformerPreTrainedModel',
   ]


if TYPE_CHECKING:
   from .configuration_autoformer import (
       AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
       AutoformerConfig,
   )

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_autoformer import (
          AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
          AutoformerForPrediction,
          AutoformerModel,
          AutoformerPreTrainedModel,
      )

else:
   import sys

   lowercase__						:      Tuple	=    _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



