




import tensorflow as tf

from ...tf_utils import shape_list




class       __lowerCAmelCase				(							tf.keras.layers.Layer):






       def __init__(      self:    Optional[Any]				,    _lowerCAmelCase:    str				,    _lowerCAmelCase:    Optional[int]				,    _lowerCAmelCase:    List[str]				,    _lowerCAmelCase:    int				,    _lowerCAmelCase:    str=1				,    _lowerCAmelCase:    Union[str, Any]=False				,    **_lowerCAmelCase:    str ):
         super().__init__(**_lowerCAmelCase )

         lowercase			:Optional[Any]							=  vocab_size
         lowercase			:Any							=  d_embed
         lowercase			:Union[str, Any]							=  d_proj

         lowercase			:Any							=  cutoffs + [vocab_size]
         lowercase			:List[str]							=  [0] + self.cutoffs
         lowercase			:Any							=  div_val

         lowercase			:List[str]							=  self.cutoffs[0]
         lowercase			:Tuple							=  len(self.cutoffs ) - 1
         lowercase			:Any							=  self.shortlist_size + self.n_clusters
         lowercase			:Any							=  keep_order

         lowercase			:Optional[Any]							=  []
         lowercase			:Any							=  []






       def     SCREAMING_SNAKE_CASE			(      self:    Union[str, Any]				,    _lowerCAmelCase:    Optional[int] ):
         if self.n_clusters > 0:
           lowercase			:str							=  self.add_weight(
               shape=(self.n_clusters, self.d_embed)				,    initializer="zeros"				,    trainable=_lowerCAmelCase				,    name="cluster_weight" )
           lowercase			:Union[str, Any]							=  self.add_weight(
               shape=(self.n_clusters,)				,    initializer="zeros"				,    trainable=_lowerCAmelCase				,    name="cluster_bias" )

         if self.div_val == 1:
           for i in range(len(self.cutoffs ) ):
             if self.d_proj != self.d_embed:
               lowercase			:Dict							=  self.add_weight(
                   shape=(self.d_embed, self.d_proj)				,    initializer="zeros"				,    trainable=_lowerCAmelCase				,    name=F"out_projs_._{i}"				,    )
               self.out_projs.append(_lowerCAmelCase )
             else:
               self.out_projs.append(_lowerCAmelCase )
             lowercase			:List[Any]							=  self.add_weight(
                 shape=(self.vocab_size, self.d_embed)				,    initializer="zeros"				,    trainable=_lowerCAmelCase				,    name=F"out_layers_._{i}_._weight"				,    )
             lowercase			:int							=  self.add_weight(
                 shape=(self.vocab_size,)				,    initializer="zeros"				,    trainable=_lowerCAmelCase				,    name=F"out_layers_._{i}_._bias"				,    )
             self.out_layers.append((weight, bias) )
         else:
           for i in range(len(self.cutoffs ) ):
             lowercase		,			lowercase			:int							=  self.cutoff_ends[i], self.cutoff_ends[i + 1]
             lowercase			:List[Any]							=  self.d_embed // (self.div_val**i)

             lowercase			:Optional[int]							=  self.add_weight(
                 shape=(d_emb_i, self.d_proj)				,    initializer="zeros"				,    trainable=_lowerCAmelCase				,    name=F"out_projs_._{i}" )
             self.out_projs.append(_lowerCAmelCase )
             lowercase			:List[str]							=  self.add_weight(
                 shape=(r_idx - l_idx, d_emb_i)				,    initializer="zeros"				,    trainable=_lowerCAmelCase				,    name=F"out_layers_._{i}_._weight"				,    )
             lowercase			:List[str]							=  self.add_weight(
                 shape=(r_idx - l_idx,)				,    initializer="zeros"				,    trainable=_lowerCAmelCase				,    name=F"out_layers_._{i}_._bias"				,    )
             self.out_layers.append((weight, bias) )
         super().build(_lowerCAmelCase )






       @staticmethod
       def     SCREAMING_SNAKE_CASE			(      _lowerCAmelCase:    Tuple				,    _lowerCAmelCase:    List[Any]				,    _lowerCAmelCase:    List[Any]				,    _lowerCAmelCase:    Dict=None ):
         lowercase			:List[str]							=  x
         if proj is not None:
           lowercase			:Dict							=  tf.einsum("ibd,ed->ibe"				,    _lowerCAmelCase				,    _lowerCAmelCase )
         return tf.einsum("ibd,nd->ibn"				,    _lowerCAmelCase				,    _lowerCAmelCase ) + b






       @staticmethod
       def     SCREAMING_SNAKE_CASE			(      _lowerCAmelCase:    Union[str, Any]				,    _lowerCAmelCase:    Tuple ):
         lowercase			:Dict							=  shape_list(_lowerCAmelCase )
         lowercase			:str							=  tf.range(lp_size[0]				,    dtype=target.dtype )
         lowercase			:Dict							=  tf.stack([r, target]				,    1 )
         return tf.gather_nd(_lowerCAmelCase				,    _lowerCAmelCase )






       def     SCREAMING_SNAKE_CASE			(      self:    List[Any]				,    _lowerCAmelCase:    List[Any]				,    _lowerCAmelCase:    List[str]				,    _lowerCAmelCase:    Optional[Any]=True				,    _lowerCAmelCase:    Any=False ):
         lowercase			:Optional[Any]							=  0
         if self.n_clusters == 0:
           lowercase			:List[str]							=  self._logit(_lowerCAmelCase				,    self.out_layers[0][0]				,    self.out_layers[0][1]				,    self.out_projs[0] )
           if target is not None:
             lowercase			:Tuple							=  tf.nn.sparse_softmax_cross_entropy_with_logits(labels=_lowerCAmelCase				,    logits=_lowerCAmelCase )
           lowercase			:Optional[Any]							=  tf.nn.log_softmax(_lowerCAmelCase				,    axis=-1 )
         else:
           lowercase			:Tuple							=  shape_list(_lowerCAmelCase )
           lowercase			:int							=  []
           lowercase			:Union[str, Any]							=  tf.zeros(hidden_sizes[:2] )
           for i in range(len(self.cutoffs ) ):
             lowercase		,			lowercase			:int							=  self.cutoff_ends[i], self.cutoff_ends[i + 1]
             if target is not None:
               lowercase			:Any							=  (target >= l_idx) & (target < r_idx)
               lowercase			:Optional[int]							=  tf.where(_lowerCAmelCase )
               lowercase			:int							=  tf.boolean_mask(_lowerCAmelCase				,    _lowerCAmelCase ) - l_idx

             if self.div_val == 1:
               lowercase			:Dict							=  self.out_layers[0][0][l_idx:r_idx]
               lowercase			:Optional[Any]							=  self.out_layers[0][1][l_idx:r_idx]
             else:
               lowercase			:Optional[Any]							=  self.out_layers[i][0]
               lowercase			:str							=  self.out_layers[i][1]

             if i == 0:
               lowercase			:Union[str, Any]							=  tf.concat([cur_W, self.cluster_weight]				,    0 )
               lowercase			:Optional[int]							=  tf.concat([cur_b, self.cluster_bias]				,    0 )

               lowercase			:List[Any]							=  self._logit(_lowerCAmelCase				,    _lowerCAmelCase				,    _lowerCAmelCase				,    self.out_projs[0] )
               lowercase			:Union[str, Any]							=  tf.nn.log_softmax(_lowerCAmelCase )
               out.append(head_logprob[..., : self.cutoffs[0]] )
               if target is not None:
                 lowercase			:Union[str, Any]							=  tf.boolean_mask(_lowerCAmelCase				,    _lowerCAmelCase )
                 lowercase			:Optional[Any]							=  self._gather_logprob(_lowerCAmelCase				,    _lowerCAmelCase )
             else:
               lowercase			:Any							=  self._logit(_lowerCAmelCase				,    _lowerCAmelCase				,    _lowerCAmelCase				,    self.out_projs[i] )
               lowercase			:Union[str, Any]							=  tf.nn.log_softmax(_lowerCAmelCase )
               lowercase			:int							=  self.cutoffs[0] + i - 1  # No probability for the head cluster
               lowercase			:Any							=  head_logprob[..., cluster_prob_idx, None] + tail_logprob
               out.append(_lowerCAmelCase )
               if target is not None:
                 lowercase			:Tuple							=  tf.boolean_mask(_lowerCAmelCase				,    _lowerCAmelCase )
                 lowercase			:Any							=  tf.boolean_mask(_lowerCAmelCase				,    _lowerCAmelCase )
                 lowercase			:int							=  self._gather_logprob(_lowerCAmelCase				,    _lowerCAmelCase )
                 cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
             if target is not None:
               loss += tf.scatter_nd(_lowerCAmelCase				,    -cur_logprob				,    shape_list(_lowerCAmelCase ) )
           lowercase			:Optional[Any]							=  tf.concat(_lowerCAmelCase				,    axis=-1 )

         if target is not None:
           if return_mean:
             lowercase			:Union[str, Any]							=  tf.reduce_mean(_lowerCAmelCase )
           # Add the training-time loss value to the layer using `self.add_loss()`.
           self.add_loss(_lowerCAmelCase )

           # Log the loss as a metric (we could log arbitrary metrics,
           # including different metrics for training and inference.
           self.add_metric(_lowerCAmelCase				,    name=self.name				,    aggregation="mean" if return_mean else "" )

         return out









from typing import TYPE_CHECKING

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


_UpperCAmelCase		:  int	=   {"configuration_swin": ["SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP", "SwinConfig", "SwinOnnxConfig"]}


try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _UpperCAmelCase		:  List[str]	=   [
         "SWIN_PRETRAINED_MODEL_ARCHIVE_LIST",
         "SwinForImageClassification",
         "SwinForMaskedImageModeling",
         "SwinModel",
         "SwinPreTrainedModel",
         "SwinBackbone",
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _UpperCAmelCase		:  int	=   [
         "TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST",
         "TFSwinForImageClassification",
         "TFSwinForMaskedImageModeling",
         "TFSwinModel",
         "TFSwinPreTrainedModel",
     ]

if TYPE_CHECKING:
     from .configuration_swin import SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP, SwinConfig, SwinOnnxConfig

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_swin import (
              SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
              SwinBackbone,
              SwinForImageClassification,
              SwinForMaskedImageModeling,
              SwinModel,
              SwinPreTrainedModel,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_swin import (
              TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFSwinForImageClassification,
              TFSwinForMaskedImageModeling,
              TFSwinModel,
              TFSwinPreTrainedModel,
          )

else:
     import sys

     _UpperCAmelCase		:  int	=   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



