




from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFSequenceClassifierOutput,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import logging
from .configuration_regnet import RegNetConfig


UpperCAmelCase_									=	logging.get_logger(__name__)

# General docstring
UpperCAmelCase_									=	'RegNetConfig'

# Base docstring
UpperCAmelCase_									=	'facebook/regnet-y-040'
UpperCAmelCase_									=	[1, 1_088, 7, 7]

# Image classification docstring
UpperCAmelCase_									=	'facebook/regnet-y-040'
UpperCAmelCase_									=	'tabby, tabby cat'

UpperCAmelCase_									=	[
    'facebook/regnet-y-040',
    # See all regnet models at https://huggingface.co/models?filter=regnet
]





class     lowerCamelCase__(					tf.keras.layers.Layer):
      def __init__(				self: Optional[Any]    ,       UpperCamelCase_: int    ,       UpperCamelCase_: int = 3    ,       UpperCamelCase_: int = 1    ,       UpperCamelCase_: int = 1    ,       UpperCamelCase_: Optional[str] = "relu"    ,       **UpperCamelCase_: List[str]    ,       ):
             super().__init__(**UpperCamelCase_   )
             # The padding and conv has been verified in
             # https://colab.research.google.com/gist/sayakpaul/854bc10eeaf21c9ee2119e0b9f3841a7/scratchpad.ipynb
             __lowerCamelCase   =						tf.keras.layers.ZeroPaddingaD(padding=kernel_size // 2   )
             __lowerCamelCase   =						tf.keras.layers.ConvaD(
                 filters=UpperCamelCase_    ,       kernel_size=UpperCamelCase_    ,       strides=UpperCamelCase_    ,       padding="""VALID"""    ,       groups=UpperCamelCase_    ,       use_bias=UpperCamelCase_    ,       name="""convolution"""    ,       )
             __lowerCamelCase   =						tf.keras.layers.BatchNormalization(epsilon=1E-5    ,       momentum=0.9    ,       name="""normalization"""   )
             __lowerCamelCase   =						ACTaFN[activation] if activation is not None else tf.identity




      def 			lowerCAmelCase__						(				self: str    ,       UpperCamelCase_: Union[str, Any]   ):
             __lowerCamelCase   =						self.convolution(self.padding(UpperCamelCase_   )   )
             __lowerCamelCase   =						self.normalization(UpperCamelCase_   )
             __lowerCamelCase   =						self.activation(UpperCamelCase_   )
             return hidden_state





class     lowerCamelCase__(					tf.keras.layers.Layer):
      def __init__(				self: Union[str, Any]    ,       UpperCamelCase_: RegNetConfig    ,       **UpperCamelCase_: Tuple   ):
             super().__init__(**UpperCamelCase_   )
             __lowerCamelCase   =						config.num_channels
             __lowerCamelCase   =						TFRegNetConvLayer(
                 out_channels=config.embedding_size    ,       kernel_size=3    ,       stride=2    ,       activation=config.hidden_act    ,       name="""embedder"""    ,       )




      def 			lowerCAmelCase__						(				self: Optional[int]    ,       UpperCamelCase_: Any   ):
             __lowerCamelCase   =						shape_list(UpperCamelCase_   )[1]
             if tf.executing_eagerly() and num_channels != self.num_channels:
                    raise ValueError(
                        """Make sure that the channel dimension of the pixel values match with the one set in the configuration."""   )

             # When running on CPU, `tf.keras.layers.Conv2D` doesn't support `NCHW` format.
             # So change the input format from `NCHW` to `NHWC`.
             # shape = (batch_size, in_height, in_width, in_channels=num_channels)
             __lowerCamelCase   =						tf.transpose(UpperCamelCase_    ,       perm=(0, 2, 3, 1)   )
             __lowerCamelCase   =						self.embedder(UpperCamelCase_   )
             return hidden_state





class     lowerCamelCase__(					tf.keras.layers.Layer):
      def __init__(				self: Optional[Any]    ,       UpperCamelCase_: int    ,       UpperCamelCase_: int = 2    ,       **UpperCamelCase_: List[Any]   ):
             super().__init__(**UpperCamelCase_   )
             __lowerCamelCase   =						tf.keras.layers.ConvaD(
                 filters=UpperCamelCase_    ,       kernel_size=1    ,       strides=UpperCamelCase_    ,       use_bias=UpperCamelCase_    ,       name="""convolution"""   )
             __lowerCamelCase   =						tf.keras.layers.BatchNormalization(epsilon=1E-5    ,       momentum=0.9    ,       name="""normalization"""   )




      def 			lowerCAmelCase__						(				self: Optional[int]    ,       UpperCamelCase_: tf.Tensor    ,       UpperCamelCase_: bool = False   ):
             return self.normalization(self.convolution(UpperCamelCase_   )    ,       training=UpperCamelCase_   )





class     lowerCamelCase__(					tf.keras.layers.Layer):
      def __init__(				self: str    ,       UpperCamelCase_: int    ,       UpperCamelCase_: int    ,       **UpperCamelCase_: str   ):
             super().__init__(**UpperCamelCase_   )
             __lowerCamelCase   =						tf.keras.layers.GlobalAveragePoolingaD(keepdims=UpperCamelCase_    ,       name="""pooler"""   )
             __lowerCamelCase   =						[
                 tf.keras.layers.ConvaD(filters=UpperCamelCase_    ,       kernel_size=1    ,       activation="""relu"""    ,       name="""attention.0"""   ),
                 tf.keras.layers.ConvaD(filters=UpperCamelCase_    ,       kernel_size=1    ,       activation="""sigmoid"""    ,       name="""attention.2"""   ),
             ]




      def 			lowerCAmelCase__						(				self: List[Any]    ,       UpperCamelCase_: str   ):
             # [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
             __lowerCamelCase   =						self.pooler(UpperCamelCase_   )
             for layer_module in self.attention:
                    __lowerCamelCase   =						layer_module(UpperCamelCase_   )
             __lowerCamelCase   =						hidden_state * pooled
             return hidden_state





class     lowerCamelCase__(					tf.keras.layers.Layer):
      def __init__(				self: Union[str, Any]    ,       UpperCamelCase_: RegNetConfig    ,       UpperCamelCase_: int    ,       UpperCamelCase_: int    ,       UpperCamelCase_: int = 1    ,       **UpperCamelCase_: Any   ):
             super().__init__(**UpperCamelCase_   )
             __lowerCamelCase   =						in_channels != out_channels or stride != 1
             __lowerCamelCase   =						max(1    ,       out_channels // config.groups_width   )
             __lowerCamelCase   =						(
                 TFRegNetShortCut(UpperCamelCase_    ,       stride=UpperCamelCase_    ,       name="""shortcut"""   )
                 if should_apply_shortcut
                 else tf.keras.layers.Activation("""linear"""    ,       name="""shortcut"""   )
             )
             # `self.layers` instead of `self.layer` because that is a reserved argument.
             __lowerCamelCase   =						[
                 TFRegNetConvLayer(UpperCamelCase_    ,       kernel_size=1    ,       activation=config.hidden_act    ,       name="""layer.0"""   ),
                 TFRegNetConvLayer(
                     UpperCamelCase_    ,       stride=UpperCamelCase_    ,       groups=UpperCamelCase_    ,       activation=config.hidden_act    ,       name="""layer.1"""   ),
                 TFRegNetConvLayer(UpperCamelCase_    ,       kernel_size=1    ,       activation=UpperCamelCase_    ,       name="""layer.2"""   ),
             ]
             __lowerCamelCase   =						ACTaFN[config.hidden_act]




      def 			lowerCAmelCase__						(				self: str    ,       UpperCamelCase_: Union[str, Any]   ):
             __lowerCamelCase   =						hidden_state
             for layer_module in self.layers:
                    __lowerCamelCase   =						layer_module(UpperCamelCase_   )
             __lowerCamelCase   =						self.shortcut(UpperCamelCase_   )
             hidden_state += residual
             __lowerCamelCase   =						self.activation(UpperCamelCase_   )
             return hidden_state





class     lowerCamelCase__(					tf.keras.layers.Layer):
      def __init__(				self: Any    ,       UpperCamelCase_: RegNetConfig    ,       UpperCamelCase_: int    ,       UpperCamelCase_: int    ,       UpperCamelCase_: int = 1    ,       **UpperCamelCase_: int   ):
             super().__init__(**UpperCamelCase_   )
             __lowerCamelCase   =						in_channels != out_channels or stride != 1
             __lowerCamelCase   =						max(1    ,       out_channels // config.groups_width   )
             __lowerCamelCase   =						(
                 TFRegNetShortCut(UpperCamelCase_    ,       stride=UpperCamelCase_    ,       name="""shortcut"""   )
                 if should_apply_shortcut
                 else tf.keras.layers.Activation("""linear"""    ,       name="""shortcut"""   )
             )
             __lowerCamelCase   =						[
                 TFRegNetConvLayer(UpperCamelCase_    ,       kernel_size=1    ,       activation=config.hidden_act    ,       name="""layer.0"""   ),
                 TFRegNetConvLayer(
                     UpperCamelCase_    ,       stride=UpperCamelCase_    ,       groups=UpperCamelCase_    ,       activation=config.hidden_act    ,       name="""layer.1"""   ),
                 TFRegNetSELayer(UpperCamelCase_    ,       reduced_channels=int(round(in_channels / 4   )   )    ,       name="""layer.2"""   ),
                 TFRegNetConvLayer(UpperCamelCase_    ,       kernel_size=1    ,       activation=UpperCamelCase_    ,       name="""layer.3"""   ),
             ]
             __lowerCamelCase   =						ACTaFN[config.hidden_act]




      def 			lowerCAmelCase__						(				self: Dict    ,       UpperCamelCase_: str   ):
             __lowerCamelCase   =						hidden_state
             for layer_module in self.layers:
                    __lowerCamelCase   =						layer_module(UpperCamelCase_   )
             __lowerCamelCase   =						self.shortcut(UpperCamelCase_   )
             hidden_state += residual
             __lowerCamelCase   =						self.activation(UpperCamelCase_   )
             return hidden_state





class     lowerCamelCase__(					tf.keras.layers.Layer):
      def __init__(				self: List[str]    ,       UpperCamelCase_: RegNetConfig    ,       UpperCamelCase_: int    ,       UpperCamelCase_: int    ,       UpperCamelCase_: int = 2    ,       UpperCamelCase_: int = 2    ,       **UpperCamelCase_: Tuple   ):
             super().__init__(**UpperCamelCase_   )

             __lowerCamelCase   =						TFRegNetXLayer if config.layer_type == """x""" else TFRegNetYLayer
             __lowerCamelCase   =						[
                 # downsampling is done in the first layer with stride of 2
                 layer(UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       stride=UpperCamelCase_    ,       name="""layers.0"""   ),
                 *[layer(UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       name=F'layers.{i+1}'   ) for i in range(depth - 1   )],
             ]




      def 			lowerCAmelCase__						(				self: str    ,       UpperCamelCase_: str   ):
             for layer_module in self.layers:
                    __lowerCamelCase   =						layer_module(UpperCamelCase_   )
             return hidden_state





class     lowerCamelCase__(					tf.keras.layers.Layer):
      def __init__(				self: List[Any]    ,       UpperCamelCase_: RegNetConfig    ,       **UpperCamelCase_: Dict   ):
             super().__init__(**UpperCamelCase_   )
             __lowerCamelCase   =						[]
             # based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
             self.stages.append(
                 TFRegNetStage(
                     UpperCamelCase_    ,       config.embedding_size    ,       config.hidden_sizes[0]    ,       stride=2 if config.downsample_in_first_stage else 1    ,       depth=config.depths[0]    ,       name="""stages.0"""    ,       )   )
             __lowerCamelCase   =						zip(config.hidden_sizes    ,       config.hidden_sizes[1:]   )
             for i, ((in_channels, out_channels), depth) in enumerate(zip(UpperCamelCase_    ,       config.depths[1:]   )   ):
                    self.stages.append(TFRegNetStage(UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       depth=UpperCamelCase_    ,       name=F'stages.{i+1}'   )   )




      def 			lowerCAmelCase__						(				self: int    ,       UpperCamelCase_: tf.Tensor    ,       UpperCamelCase_: bool = False    ,       UpperCamelCase_: bool = True   ):
             __lowerCamelCase   =						() if output_hidden_states else None

             for stage_module in self.stages:
                    if output_hidden_states:
                           __lowerCamelCase   =						hidden_states + (hidden_state,)

                    __lowerCamelCase   =						stage_module(UpperCamelCase_   )

             if output_hidden_states:
                    __lowerCamelCase   =						hidden_states + (hidden_state,)

             if not return_dict:
                    return tuple(v for v in [hidden_state, hidden_states] if v is not None   )

             return TFBaseModelOutputWithNoAttention(last_hidden_state=UpperCamelCase_    ,       hidden_states=UpperCamelCase_   )





@keras_serializable
class     lowerCamelCase__(					tf.keras.layers.Layer):
      UpperCAmelCase__      :	Union[str, Any]    		=   RegNetConfig
      def __init__(				self: str    ,       UpperCamelCase_: str    ,       **UpperCamelCase_: List[str]   ):
             super().__init__(**UpperCamelCase_   )
             __lowerCamelCase   =						config
             __lowerCamelCase   =						TFRegNetEmbeddings(UpperCamelCase_    ,       name="""embedder"""   )
             __lowerCamelCase   =						TFRegNetEncoder(UpperCamelCase_    ,       name="""encoder"""   )
             __lowerCamelCase   =						tf.keras.layers.GlobalAveragePoolingaD(keepdims=UpperCamelCase_    ,       name="""pooler"""   )




      @unpack_inputs
      def 			lowerCAmelCase__						(				self: Optional[int]    ,       UpperCamelCase_: tf.Tensor    ,       UpperCamelCase_: Optional[bool] = None    ,       UpperCamelCase_: Optional[bool] = None    ,       UpperCamelCase_: bool = False    ,       ):
             __lowerCamelCase   =						(
                 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
             )
             __lowerCamelCase   =						return_dict if return_dict is not None else self.config.use_return_dict

             __lowerCamelCase   =						self.embedder(UpperCamelCase_    ,       training=UpperCamelCase_   )

             __lowerCamelCase   =						self.encoder(
                 UpperCamelCase_    ,       output_hidden_states=UpperCamelCase_    ,       return_dict=UpperCamelCase_    ,       training=UpperCamelCase_   )

             __lowerCamelCase   =						encoder_outputs[0]
             __lowerCamelCase   =						self.pooler(UpperCamelCase_   )

             # Change to NCHW output format have uniformity in the modules
             __lowerCamelCase   =						tf.transpose(UpperCamelCase_    ,       perm=(0, 3, 1, 2)   )
             __lowerCamelCase   =						tf.transpose(UpperCamelCase_    ,       perm=(0, 3, 1, 2)   )

             # Change the other hidden state outputs to NCHW as well
             if output_hidden_states:
                    __lowerCamelCase   =						tuple([tf.transpose(UpperCamelCase_    ,       perm=(0, 3, 1, 2)   ) for h in encoder_outputs[1]]   )

             if not return_dict:
                    return (last_hidden_state, pooled_output) + encoder_outputs[1:]

             return TFBaseModelOutputWithPoolingAndNoAttention(
                 last_hidden_state=UpperCamelCase_    ,       pooler_output=UpperCamelCase_    ,       hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states    ,       )





class     lowerCamelCase__(					__lowerCamelCase):
      UpperCAmelCase__      :	Optional[Any]    		=   RegNetConfig
      UpperCAmelCase__      :	str    		=   'regnet'
      UpperCAmelCase__      :	Union[str, Any]    		=   'pixel_values'
      @property
      def 			lowerCAmelCase__						(				self: Tuple   ):
             return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 2_24, 2_24)    ,       dtype=tf.floataa   )}



UpperCAmelCase_									=	r'\n    Parameters:\n    This model is a Tensorflow\n    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a\n    regular Tensorflow Module and refer to the Tensorflow documentation for all matter related to general usage and\n    behavior.\n        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.\n'

UpperCAmelCase_									=	r'\n    Args:\n        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConveNextImageProcessor.__call__`] for details.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n'





@add_start_docstrings(
    'The bare RegNet model outputting raw features without any specific head on top.' ,     __lowerCamelCase ,     )
class     lowerCamelCase__(					__lowerCamelCase):
      def __init__(				self: Optional[Any]    ,       UpperCamelCase_: RegNetConfig    ,       *UpperCamelCase_: List[str]    ,       **UpperCamelCase_: Any   ):
             super().__init__(UpperCamelCase_    ,       *UpperCamelCase_    ,       **UpperCamelCase_   )
             __lowerCamelCase   =						TFRegNetMainLayer(UpperCamelCase_    ,       name="""regnet"""   )




      @unpack_inputs
      @add_start_docstrings_to_model_forward(UpperCamelCase_   )
      @add_code_sample_docstrings(
          checkpoint=_CHECKPOINT_FOR_DOC    ,       output_type=UpperCamelCase_    ,       config_class=_CONFIG_FOR_DOC    ,       modality="""vision"""    ,       expected_output=_EXPECTED_OUTPUT_SHAPE    ,       )
      def 			lowerCAmelCase__						(				self: Tuple    ,       UpperCamelCase_: tf.Tensor    ,       UpperCamelCase_: Optional[bool] = None    ,       UpperCamelCase_: Optional[bool] = None    ,       UpperCamelCase_: int=False    ,       ):
             __lowerCamelCase   =						(
                 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
             )
             __lowerCamelCase   =						return_dict if return_dict is not None else self.config.use_return_dict

             __lowerCamelCase   =						self.regnet(
                 pixel_values=UpperCamelCase_    ,       output_hidden_states=UpperCamelCase_    ,       return_dict=UpperCamelCase_    ,       training=UpperCamelCase_    ,       )
             if not return_dict:
                    return (outputs[0],) + outputs[1:]

             return TFBaseModelOutputWithPoolingAndNoAttention(
                 last_hidden_state=outputs.last_hidden_state    ,       pooler_output=outputs.pooler_output    ,       hidden_states=outputs.hidden_states    ,       )





@add_start_docstrings(
    '\n    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    ' ,     __lowerCamelCase ,     )
class     lowerCamelCase__(					__lowerCamelCase ,     __lowerCamelCase):
      def __init__(				self: Optional[int]    ,       UpperCamelCase_: RegNetConfig    ,       *UpperCamelCase_: Union[str, Any]    ,       **UpperCamelCase_: Dict   ):
             super().__init__(UpperCamelCase_    ,       *UpperCamelCase_    ,       **UpperCamelCase_   )
             __lowerCamelCase   =						config.num_labels
             __lowerCamelCase   =						TFRegNetMainLayer(UpperCamelCase_    ,       name="""regnet"""   )
             # classification head
             __lowerCamelCase   =						[
                 tf.keras.layers.Flatten(),
                 tf.keras.layers.Dense(config.num_labels    ,       name="""classifier.1"""   ) if config.num_labels > 0 else tf.identity,
             ]




      @unpack_inputs
      @add_start_docstrings_to_model_forward(UpperCamelCase_   )
      @add_code_sample_docstrings(
          checkpoint=_IMAGE_CLASS_CHECKPOINT    ,       output_type=UpperCamelCase_    ,       config_class=_CONFIG_FOR_DOC    ,       expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT    ,       )
      def 			lowerCAmelCase__						(				self: Any    ,       UpperCamelCase_: tf.Tensor = None    ,       UpperCamelCase_: tf.Tensor = None    ,       UpperCamelCase_: bool = None    ,       UpperCamelCase_: bool = None    ,       UpperCamelCase_: List[str]=False    ,       ):
             __lowerCamelCase   =						(
                 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
             )
             __lowerCamelCase   =						return_dict if return_dict is not None else self.config.use_return_dict

             __lowerCamelCase   =						self.regnet(
                 UpperCamelCase_    ,       output_hidden_states=UpperCamelCase_    ,       return_dict=UpperCamelCase_    ,       training=UpperCamelCase_   )

             __lowerCamelCase   =						outputs.pooler_output if return_dict else outputs[1]

             __lowerCamelCase   =						self.classifier[0](UpperCamelCase_   )
             __lowerCamelCase   =						self.classifier[1](UpperCamelCase_   )

             __lowerCamelCase   =						None if labels is None else self.hf_compute_loss(labels=UpperCamelCase_    ,       logits=UpperCamelCase_   )

             if not return_dict:
                    __lowerCamelCase   =						(logits,) + outputs[2:]
                    return ((loss,) + output) if loss is not None else output

             return TFSequenceClassifierOutput(loss=UpperCamelCase_    ,       logits=UpperCamelCase_    ,       hidden_states=outputs.hidden_states   )




'''simple docstring'''






import torch
from torch import nn






class 			lowercase_	(nn.Module      ):




  """simple docstring"""




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

   __lowercase							  = n_token
   __lowercase							  = d_embed
   __lowercase							  = d_proj

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

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

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

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

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

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

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

   __lowercase							  = keep_order






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]					,lowercase__   :     List[Any]					,lowercase__   :     Dict					,lowercase__   :     List[Any]					,lowercase__   :     Any			):
   if proj is None:
    __lowercase							  = nn.functional.linear(lowercase__					,lowercase__					,bias=lowercase__			)
   else:
    # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
    __lowercase							  = nn.functional.linear(lowercase__					,proj.t().contiguous()			)
    __lowercase							  = nn.functional.linear(lowercase__					,lowercase__					,bias=lowercase__			)
    # else:
    #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
    #     if bias is not None:
    #         logit = logit + bias

   return logit






  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]					,lowercase__   :     Optional[Any]					,lowercase__   :     Any=None					,lowercase__   :     List[str]=False			):

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

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

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

     weights.append(lowercase__			)
     biases.append(lowercase__			)

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

    __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)
    __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=1			)

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

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

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

      if indices_i.numel() == 0:
       continue

      __lowercase							  = labels.index_select(0					,lowercase__			) - l_idx
      __lowercase							  = head_logprob.index_select(0					,lowercase__			)
      __lowercase							  = hidden.index_select(0					,lowercase__			)
     else:
      __lowercase							  = hidden

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

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

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

   return out






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

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

     weights.append(lowercase__			)
     biases.append(lowercase__			)

    __lowercase  ,	__lowercase  ,	__lowercase							  = weights[0], biases[0], self.out_projs[0]
    __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)

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

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

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

      __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)
      __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=1			)

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

    return out

