





import torch
from torch import nn



class 				_UpperCAmelCase    (    nn.Module    ):

  """simple docstring"""


  def __init__(				self :		int					,  lowerCAmelCase_ :		str					,  lowerCAmelCase_ :		Dict					,  lowerCAmelCase_ :		Dict					,  lowerCAmelCase_ :		List[str]					,  lowerCAmelCase_ :		Optional[Any]=1					,  lowerCAmelCase_ :		Any=False     )							->						Dict:
     super().__init__()

     __lowerCAmelCase								=	n_token
     __lowerCAmelCase								=	d_embed
     __lowerCAmelCase								=	d_proj

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

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

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

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

     __lowerCAmelCase								=	keep_order






  def 				lowercase		(				self :		Tuple					,  lowerCAmelCase_ :		Optional[Any]					,  lowerCAmelCase_ :		str					,  lowerCAmelCase_ :		Optional[Any]					,  lowerCAmelCase_ :		List[Any]     )							->						List[str]:
     if proj is None:
        __lowerCAmelCase								=	nn.functional.linear(lowerCAmelCase_					,  lowerCAmelCase_					,  bias=lowerCAmelCase_     )
     else:
        # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
        __lowerCAmelCase								=	nn.functional.linear(lowerCAmelCase_					,  proj.t().contiguous()     )
        __lowerCAmelCase								=	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 				lowercase		(				self :		Optional[Any]					,  lowerCAmelCase_ :		Optional[Any]					,  lowerCAmelCase_ :		int=None					,  lowerCAmelCase_ :		Any=False     )							->						Optional[int]:

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

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

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

           weights.append(lowerCAmelCase_     )
           biases.append(lowerCAmelCase_     )

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

        __lowerCAmelCase								=	self._compute_logit(lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_     )
        __lowerCAmelCase								=	nn.functional.log_softmax(lowerCAmelCase_					,  dim=1     )

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

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

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

              if indices_i.numel() == 0:
                 continue

              __lowerCAmelCase								=	labels.index_select(0					,  lowerCAmelCase_     ) - l_idx
              __lowerCAmelCase								=	head_logprob.index_select(0					,  lowerCAmelCase_     )
              __lowerCAmelCase								=	hidden.index_select(0					,  lowerCAmelCase_     )
           else:
              __lowerCAmelCase								=	hidden

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

              __lowerCAmelCase								=	self._compute_logit(lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_     )
              __lowerCAmelCase								=	nn.functional.log_softmax(lowerCAmelCase_					,  dim=1     )
              __lowerCAmelCase								=	self.cutoffs[0] + i - 1  # No probability for the head cluster
              if labels is not None:
                 __lowerCAmelCase								=	head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
                     1					,  target_i[:, None]     ).squeeze(1     )
              else:
                 __lowerCAmelCase								=	head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
                 __lowerCAmelCase								=	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 				lowercase		(				self :		Optional[int]					,  lowerCAmelCase_ :		List[str]     )							->						Any:
     if self.n_clusters == 0:
        __lowerCAmelCase								=	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
        __lowerCAmelCase	,	__lowerCAmelCase								=	[], []
        for i in range(len(self.cutoffs     )     ):
           if self.div_val == 1:
              __lowerCAmelCase	,	__lowerCAmelCase								=	self.cutoff_ends[i], self.cutoff_ends[i + 1]
              __lowerCAmelCase								=	self.out_layers[0].weight[l_idx:r_idx]
              __lowerCAmelCase								=	self.out_layers[0].bias[l_idx:r_idx]
           else:
              __lowerCAmelCase								=	self.out_layers[i].weight
              __lowerCAmelCase								=	self.out_layers[i].bias

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

           weights.append(lowerCAmelCase_     )
           biases.append(lowerCAmelCase_     )

        __lowerCAmelCase	,	__lowerCAmelCase	,	__lowerCAmelCase								=	weights[0], biases[0], self.out_projs[0]
        __lowerCAmelCase								=	self._compute_logit(lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_     )

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

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

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

              __lowerCAmelCase								=	self._compute_logit(lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_     )
              __lowerCAmelCase								=	nn.functional.log_softmax(lowerCAmelCase_					,  dim=1     )

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

        return out








from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_snake_case  :    List[Any]						    =  logging.get_logger(__name__)

_snake_case  :    List[Any]						    =  {
    'microsoft/beit-base-patch16-224-pt22k': (
        'https://huggingface.co/microsoft/beit-base-patch16-224-pt22k/resolve/main/config.json'
    ),
    # See all BEiT models at https://huggingface.co/models?filter=beit
}



class 				_UpperCAmelCase    (    _UpperCamelCase    ):

  """simple docstring"""


  a_        = """beit"""






  def __init__(				self :		List[Any]					,  lowerCAmelCase_ :		Tuple=8_1_9_2					,  lowerCAmelCase_ :		Optional[int]=7_6_8					,  lowerCAmelCase_ :		int=1_2					,  lowerCAmelCase_ :		Optional[int]=1_2					,  lowerCAmelCase_ :		Any=3_0_7_2					,  lowerCAmelCase_ :		Optional[int]="gelu"					,  lowerCAmelCase_ :		Any=0.0					,  lowerCAmelCase_ :		Any=0.0					,  lowerCAmelCase_ :		Any=0.02					,  lowerCAmelCase_ :		int=1e-12					,  lowerCAmelCase_ :		int=2_2_4					,  lowerCAmelCase_ :		str=1_6					,  lowerCAmelCase_ :		int=3					,  lowerCAmelCase_ :		Dict=False					,  lowerCAmelCase_ :		int=False					,  lowerCAmelCase_ :		List[Any]=False					,  lowerCAmelCase_ :		int=False					,  lowerCAmelCase_ :		List[str]=0.1					,  lowerCAmelCase_ :		Union[str, Any]=0.1					,  lowerCAmelCase_ :		List[str]=True					,  lowerCAmelCase_ :		List[Any]=[3, 5, 7, 1_1]					,  lowerCAmelCase_ :		Optional[Any]=[1, 2, 3, 6]					,  lowerCAmelCase_ :		Tuple=True					,  lowerCAmelCase_ :		Dict=0.4					,  lowerCAmelCase_ :		Tuple=2_5_6					,  lowerCAmelCase_ :		Any=1					,  lowerCAmelCase_ :		Any=False					,  lowerCAmelCase_ :		Optional[int]=2_5_5					,  **lowerCAmelCase_ :		Any					,  )							->						Dict:
     super().__init__(**lowerCAmelCase_     )

     __lowerCAmelCase								=	vocab_size
     __lowerCAmelCase								=	hidden_size
     __lowerCAmelCase								=	num_hidden_layers
     __lowerCAmelCase								=	num_attention_heads
     __lowerCAmelCase								=	intermediate_size
     __lowerCAmelCase								=	hidden_act
     __lowerCAmelCase								=	hidden_dropout_prob
     __lowerCAmelCase								=	attention_probs_dropout_prob
     __lowerCAmelCase								=	initializer_range
     __lowerCAmelCase								=	layer_norm_eps

     __lowerCAmelCase								=	image_size
     __lowerCAmelCase								=	patch_size
     __lowerCAmelCase								=	num_channels
     __lowerCAmelCase								=	use_mask_token
     __lowerCAmelCase								=	use_absolute_position_embeddings
     __lowerCAmelCase								=	use_relative_position_bias
     __lowerCAmelCase								=	use_shared_relative_position_bias
     __lowerCAmelCase								=	layer_scale_init_value
     __lowerCAmelCase								=	drop_path_rate
     __lowerCAmelCase								=	use_mean_pooling
     # decode head attributes (semantic segmentation)
     __lowerCAmelCase								=	out_indices
     __lowerCAmelCase								=	pool_scales
     # auxiliary head attributes (semantic segmentation)
     __lowerCAmelCase								=	use_auxiliary_head
     __lowerCAmelCase								=	auxiliary_loss_weight
     __lowerCAmelCase								=	auxiliary_channels
     __lowerCAmelCase								=	auxiliary_num_convs
     __lowerCAmelCase								=	auxiliary_concat_input
     __lowerCAmelCase								=	semantic_loss_ignore_index





class 				_UpperCAmelCase    (    _UpperCamelCase    ):

  """simple docstring"""


  a_        = version.parse("""1.11"""    )






  @property
  def 				lowercase		(				self :		Tuple     )							->						Mapping[str, Mapping[int, str]]:
     return OrderedDict(
         [
             ('pixel_values', {0: 'batch', 1: 'num_channels', 2: 'height', 3: 'width'}),
         ]     )






  @property
  def 				lowercase		(				self :		Optional[Any]     )							->						float:
     return 1e-4

