code: 
  
  
  from typing import Optional, Tuple, Union
  
  import tensorflow as tf
  
  from ...activations_tf import ACTaFN
  from ...modeling_tf_outputs import (
      TFBaseModelOutputWithNoAttention,
      TFBaseModelOutputWithPoolingAndNoAttention,
      TFImageClassifierOutputWithNoAttention,
  )
  from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
  from ...tf_utils import shape_list
  from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
  from .configuration_resnet import ResNetConfig
  
  
  SCREAMING_SNAKE_CASE__				:    List[Any]         =    logging.get_logger(__name__)
  
  # General docstring
  SCREAMING_SNAKE_CASE__				:    Optional[int]         =    'ResNetConfig'
  
  # Base docstring
  SCREAMING_SNAKE_CASE__				:    Any         =    'microsoft/resnet-50'
  SCREAMING_SNAKE_CASE__				:    str         =    [1, 2_048, 7, 7]
  
  # Image classification docstring
  SCREAMING_SNAKE_CASE__				:    Any         =    'microsoft/resnet-50'
  SCREAMING_SNAKE_CASE__				:    str         =    'tiger cat'
  
  SCREAMING_SNAKE_CASE__				:    List[Any]         =    [
      'microsoft/resnet-50',
      # See all resnet models at https://huggingface.co/models?filter=resnet
  ]
  
  
  class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
  
  
       def __init__(     self       ,				A_       ,				A_ = 3       ,				A_ = 1       ,				A_ = "relu"       ,				**A_						):
         super().__init__(**A_						)
         _UpperCAmelCase		:       Union[str, Any]       =    kernel_size // 2
         _UpperCAmelCase		:       Tuple       =    tf.keras.layers.ConvaD(
             A_       ,				kernel_size=A_       ,				strides=A_       ,				padding="""valid"""       ,				use_bias=A_       ,				name="""convolution"""						)
         # Use same default momentum and epsilon as PyTorch equivalent
         _UpperCAmelCase		:       Optional[int]       =    tf.keras.layers.BatchNormalization(epsilon=1e-5       ,				momentum=0.9       ,				name="""normalization"""						)
         _UpperCAmelCase		:       Tuple       =    ACTaFN[activation] if activation is not None else tf.keras.layers.Activation("""linear"""						)
  
  
       def 					__snake_case(     self       ,				A_						):
         # Pad to match that done in the PyTorch Conv2D model
         _UpperCAmelCase		:       Any       =    (self.pad_value, self.pad_value)
         _UpperCAmelCase		:       Tuple       =    tf.pad(A_       ,				[(0, 0), height_pad, width_pad, (0, 0)]						)
         _UpperCAmelCase		:       List[Any]       =    self.conv(A_						)
         return hidden_state
  
  
  
  
       def 					__snake_case(     self       ,				A_       ,				A_ = False						):
         _UpperCAmelCase		:       Any       =    self.convolution(A_						)
         _UpperCAmelCase		:       Optional[Any]       =    self.normalization(A_       ,				training=A_						)
         _UpperCAmelCase		:       List[str]       =    self.activation(A_						)
         return hidden_state
  
  
  class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
  
  
       def __init__(     self       ,				A_       ,				**A_						):
         super().__init__(**A_						)
         _UpperCAmelCase		:       str       =    TFResNetConvLayer(
             config.embedding_size       ,				kernel_size=7       ,				stride=2       ,				activation=config.hidden_act       ,				name="""embedder"""       ,				)
         _UpperCAmelCase		:       Any       =    tf.keras.layers.MaxPoolaD(pool_size=3       ,				strides=2       ,				padding="""valid"""       ,				name="""pooler"""						)
         _UpperCAmelCase		:       Dict       =    config.num_channels
  
  
  
  
       def 					__snake_case(     self       ,				A_       ,				A_ = False						):
         _UpperCAmelCase,_UpperCAmelCase,_UpperCAmelCase,_UpperCAmelCase		:       Optional[int]       =    shape_list(A_						)
         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."""						)
         _UpperCAmelCase		:       Optional[Any]       =    pixel_values
         _UpperCAmelCase		:       List[str]       =    self.embedder(A_						)
         _UpperCAmelCase		:       str       =    tf.pad(A_       ,				[[0, 0], [1, 1], [1, 1], [0, 0]]						)
         _UpperCAmelCase		:       List[Any]       =    self.pooler(A_						)
         return hidden_state
  
  
  class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
  
  
       def __init__(     self       ,				A_       ,				A_ = 2       ,				**A_						):
         super().__init__(**A_						)
         _UpperCAmelCase		:       Optional[int]       =    tf.keras.layers.ConvaD(
             A_       ,				kernel_size=1       ,				strides=A_       ,				use_bias=A_       ,				name="""convolution"""						)
         # Use same default momentum and epsilon as PyTorch equivalent
         _UpperCAmelCase		:       List[str]       =    tf.keras.layers.BatchNormalization(epsilon=1e-5       ,				momentum=0.9       ,				name="""normalization"""						)
  
  
  
  
       def 					__snake_case(     self       ,				A_       ,				A_ = False						):
         _UpperCAmelCase		:       Union[str, Any]       =    x
         _UpperCAmelCase		:       Optional[int]       =    self.convolution(A_						)
         _UpperCAmelCase		:       int       =    self.normalization(A_       ,				training=A_						)
         return hidden_state
  
  
  class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
  
  
       def __init__(     self       ,				A_       ,				A_       ,				A_ = 1       ,				A_ = "relu"       ,				**A_						):
         super().__init__(**A_						)
         _UpperCAmelCase		:       Optional[int]       =    in_channels != out_channels or stride != 1
         _UpperCAmelCase		:       str       =    TFResNetConvLayer(A_       ,				stride=A_       ,				name="""layer.0"""						)
         _UpperCAmelCase		:       Optional[Any]       =    TFResNetConvLayer(A_       ,				activation=A_       ,				name="""layer.1"""						)
         _UpperCAmelCase		:       Union[str, Any]       =    (
             TFResNetShortCut(A_       ,				stride=A_       ,				name="""shortcut"""						)
             if should_apply_shortcut
             else tf.keras.layers.Activation("""linear"""       ,				name="""shortcut"""						)
         )
         _UpperCAmelCase		:       Tuple       =    ACTaFN[activation]
  
  
  
  
       def 					__snake_case(     self       ,				A_       ,				A_ = False						):
         _UpperCAmelCase		:       Dict       =    hidden_state
         _UpperCAmelCase		:       Union[str, Any]       =    self.conva(A_       ,				training=A_						)
         _UpperCAmelCase		:       Union[str, Any]       =    self.conva(A_       ,				training=A_						)
         _UpperCAmelCase		:       int       =    self.shortcut(A_       ,				training=A_						)
         hidden_state += residual
         _UpperCAmelCase		:       List[str]       =    self.activation(A_						)
         return hidden_state
  
  
  class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
  
  
       def __init__(     self       ,				A_       ,				A_       ,				A_ = 1       ,				A_ = "relu"       ,				A_ = 4       ,				**A_       ,				):
         super().__init__(**A_						)
         _UpperCAmelCase		:       Any       =    in_channels != out_channels or stride != 1
         _UpperCAmelCase		:       List[str]       =    out_channels // reduction
         _UpperCAmelCase		:       Any       =    TFResNetConvLayer(A_       ,				kernel_size=1       ,				name="""layer.0"""						)
         _UpperCAmelCase		:       List[Any]       =    TFResNetConvLayer(A_       ,				stride=A_       ,				name="""layer.1"""						)
         _UpperCAmelCase		:       Tuple       =    TFResNetConvLayer(A_       ,				kernel_size=1       ,				activation=A_       ,				name="""layer.2"""						)
         _UpperCAmelCase		:       int       =    (
             TFResNetShortCut(A_       ,				stride=A_       ,				name="""shortcut"""						)
             if should_apply_shortcut
             else tf.keras.layers.Activation("""linear"""       ,				name="""shortcut"""						)
         )
         _UpperCAmelCase		:       int       =    ACTaFN[activation]
  
  
  
  
       def 					__snake_case(     self       ,				A_       ,				A_ = False						):
         _UpperCAmelCase		:       Tuple       =    hidden_state
         _UpperCAmelCase		:       str       =    self.conva(A_       ,				training=A_						)
         _UpperCAmelCase		:       List[Any]       =    self.conva(A_       ,				training=A_						)
         _UpperCAmelCase		:       int       =    self.conva(A_       ,				training=A_						)
         _UpperCAmelCase		:       List[Any]       =    self.shortcut(A_       ,				training=A_						)
         hidden_state += residual
         _UpperCAmelCase		:       str       =    self.activation(A_						)
         return hidden_state
  
  
  class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
  
  
       def __init__(     self       ,				A_       ,				A_       ,				A_       ,				A_ = 2       ,				A_ = 2       ,				**A_						):
         super().__init__(**A_						)
  
         _UpperCAmelCase		:       Optional[int]       =    TFResNetBottleNeckLayer if config.layer_type == """bottleneck""" else TFResNetBasicLayer
  
         _UpperCAmelCase		:       Tuple       =    [layer(A_       ,				A_       ,				stride=A_       ,				activation=config.hidden_act       ,				name="""layers.0"""						)]
         layers += [
             layer(A_       ,				A_       ,				activation=config.hidden_act       ,				name=F'''layers.{i + 1}'''						)
             for i in range(depth - 1						)
         ]
         _UpperCAmelCase		:       Any       =    layers
  
  
  
  
       def 					__snake_case(     self       ,				A_       ,				A_ = False						):
         for layer in self.stage_layers:
           _UpperCAmelCase		:       Union[str, Any]       =    layer(A_       ,				training=A_						)
         return hidden_state
  
  
  class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
  
  
       def __init__(     self       ,				A_       ,				**A_						):
         super().__init__(**A_						)
         # based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
         _UpperCAmelCase		:       Dict       =    [
             TFResNetStage(
                 A_       ,				config.embedding_size       ,				config.hidden_sizes[0]       ,				stride=2 if config.downsample_in_first_stage else 1       ,				depth=config.depths[0]       ,				name="""stages.0"""       ,				)
         ]
         for i, (in_channels, out_channels, depth) in enumerate(
             zip(config.hidden_sizes       ,				config.hidden_sizes[1:]       ,				config.depths[1:]						)						):
           self.stages.append(TFResNetStage(A_       ,				A_       ,				A_       ,				depth=A_       ,				name=F'''stages.{i + 1}'''						)						)
  
  
  
  
       def 					__snake_case(     self       ,				A_       ,				A_ = False       ,				A_ = True       ,				A_ = False       ,				):
         _UpperCAmelCase		:       Dict       =    () if output_hidden_states else None
  
         for stage_module in self.stages:
           if output_hidden_states:
             _UpperCAmelCase		:       int       =    hidden_states + (hidden_state,)
  
           _UpperCAmelCase		:       List[Any]       =    stage_module(A_       ,				training=A_						)
  
         if output_hidden_states:
           _UpperCAmelCase		:       Dict       =    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=A_       ,				hidden_states=A_						)
  
  
  class 		_SCREAMING_SNAKE_CASE    (	A    ):
       __SCREAMING_SNAKE_CASE    				=						ResNetConfig
       __SCREAMING_SNAKE_CASE    				=						'''resnet'''
       __SCREAMING_SNAKE_CASE    				=						'''pixel_values'''
  
  
       @property
       def 					__snake_case(     self						):
         return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 2_24, 2_24)       ,				dtype=tf.floataa						)}
  
  
  SCREAMING_SNAKE_CASE__				:    List[Any]         =    R'\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\n    Parameters:\n        config ([`ResNetConfig`]): 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'
  
  
  SCREAMING_SNAKE_CASE__				:    str         =    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            [`ConvNextImageProcessor.__call__`] for details.\n\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'
  
  
  @keras_serializable
  class 		_SCREAMING_SNAKE_CASE    (	tf.keras.layers.Layer    ):
       __SCREAMING_SNAKE_CASE    				=						ResNetConfig
  
  
       def __init__(     self       ,				A_       ,				**A_						):
         super().__init__(**A_						)
         _UpperCAmelCase		:       int       =    config
         _UpperCAmelCase		:       str       =    TFResNetEmbeddings(A_       ,				name="""embedder"""						)
         _UpperCAmelCase		:       int       =    TFResNetEncoder(A_       ,				name="""encoder"""						)
         _UpperCAmelCase		:       List[Any]       =    tf.keras.layers.GlobalAveragePoolingaD(keepdims=A_						)
  
  
  
  
       @unpack_inputs
       def 					__snake_case(     self       ,				A_       ,				A_ = None       ,				A_ = None       ,				A_ = False       ,				):
         _UpperCAmelCase		:       Optional[Any]       =    (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         _UpperCAmelCase		:       int       =    return_dict if return_dict is not None else self.config.use_return_dict
  
         # TF 2.0 image layers can't use NCHW format when running on CPU.
         # We transpose to NHWC format and then transpose back after the full forward pass.
         # (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels)
         _UpperCAmelCase		:       Any       =    tf.transpose(A_       ,				perm=[0, 2, 3, 1]						)
         _UpperCAmelCase		:       Union[str, Any]       =    self.embedder(A_       ,				training=A_						)
  
         _UpperCAmelCase		:       Tuple       =    self.encoder(
             A_       ,				output_hidden_states=A_       ,				return_dict=A_       ,				training=A_						)
  
         _UpperCAmelCase		:       Any       =    encoder_outputs[0]
  
         _UpperCAmelCase		:       Any       =    self.pooler(A_						)
  
         # Transpose all the outputs to the NCHW format
         # (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
         _UpperCAmelCase		:       Any       =    tf.transpose(A_       ,				(0, 3, 1, 2)						)
         _UpperCAmelCase		:       Dict       =    tf.transpose(A_       ,				(0, 3, 1, 2)						)
         _UpperCAmelCase		:       int       =    ()
         for hidden_state in encoder_outputs[1:]:
           _UpperCAmelCase		:       Tuple       =    hidden_states + tuple(tf.transpose(A_       ,				(0, 3, 1, 2)						) for h in hidden_state						)
  
         if not return_dict:
           return (last_hidden_state, pooled_output) + hidden_states
  
         _UpperCAmelCase		:       Tuple       =    hidden_states if output_hidden_states else None
  
         return TFBaseModelOutputWithPoolingAndNoAttention(
             last_hidden_state=A_       ,				pooler_output=A_       ,				hidden_states=A_       ,				)
  
  
  @add_start_docstrings(
      '''The bare ResNet model outputting raw features without any specific head on top.'''       ,  A       ,  )
  class 		_SCREAMING_SNAKE_CASE    (	A    ):
  
  
       def __init__(     self       ,				A_       ,				**A_						):
         super().__init__(A_       ,				**A_						)
         _UpperCAmelCase		:       int       =    TFResNetMainLayer(config=A_       ,				name="""resnet"""						)
  
  
  
  
       @add_start_docstrings_to_model_forward(A_						)
       @add_code_sample_docstrings(
           checkpoint=_CHECKPOINT_FOR_DOC       ,				output_type=A_       ,				config_class=_CONFIG_FOR_DOC       ,				modality="""vision"""       ,				expected_output=_EXPECTED_OUTPUT_SHAPE       ,				)
       @unpack_inputs
       def 					__snake_case(     self       ,				A_       ,				A_ = None       ,				A_ = None       ,				A_ = False       ,				):
         _UpperCAmelCase		:       int       =    (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         _UpperCAmelCase		:       Union[str, Any]       =    return_dict if return_dict is not None else self.config.use_return_dict
  
         _UpperCAmelCase		:       List[str]       =    self.resnet(
             pixel_values=A_       ,				output_hidden_states=A_       ,				return_dict=A_       ,				training=A_       ,				)
         return resnet_outputs
  
  
  @add_start_docstrings(
      '''
      ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
      ImageNet.
      '''       ,  A       ,  )
  class 		_SCREAMING_SNAKE_CASE    (	A       ,  A    ):
  
  
       def __init__(     self       ,				A_       ,				**A_						):
         super().__init__(A_       ,				**A_						)
         _UpperCAmelCase		:       int       =    config.num_labels
         _UpperCAmelCase		:       str       =    TFResNetMainLayer(A_       ,				name="""resnet"""						)
         # classification head
         _UpperCAmelCase		:       str       =    (
             tf.keras.layers.Dense(config.num_labels       ,				name="""classifier.1"""						)
             if config.num_labels > 0
             else tf.keras.layers.Activation("""linear"""       ,				name="""classifier.1"""						)
         )
  
  
       def 					__snake_case(     self       ,				A_						):
         _UpperCAmelCase		:       int       =    tf.keras.layers.Flatten()(A_						)
         _UpperCAmelCase		:       str       =    self.classifier_layer(A_						)
         return logits
  
  
  
  
       @add_start_docstrings_to_model_forward(A_						)
       @add_code_sample_docstrings(
           checkpoint=_IMAGE_CLASS_CHECKPOINT       ,				output_type=A_       ,				config_class=_CONFIG_FOR_DOC       ,				expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT       ,				)
       @unpack_inputs
       def 					__snake_case(     self       ,				A_ = None       ,				A_ = None       ,				A_ = None       ,				A_ = None       ,				A_ = False       ,				):
         _UpperCAmelCase		:       List[Any]       =    return_dict if return_dict is not None else self.config.use_return_dict
  
         _UpperCAmelCase		:       List[Any]       =    self.resnet(
             A_       ,				output_hidden_states=A_       ,				return_dict=A_       ,				training=A_						)
  
         _UpperCAmelCase		:       Optional[int]       =    outputs.pooler_output if return_dict else outputs[1]
  
         _UpperCAmelCase		:       str       =    self.classifier(A_						)
  
         _UpperCAmelCase		:       Optional[Any]       =    None if labels is None else self.hf_compute_loss(A_       ,				A_						)
  
         if not return_dict:
           _UpperCAmelCase		:       Optional[int]       =    (logits,) + outputs[2:]
           return (loss,) + output if loss is not None else output
  
         return TFImageClassifierOutputWithNoAttention(loss=A_       ,				logits=A_       ,				hidden_states=outputs.hidden_states						)
  
  
  
  
  
  
style_context: 
  
  
  import logging
  
  import torch
  from torch import nn
  from torch.nn import CrossEntropyLoss, MSELoss
  
  from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
  from transformers.models.albert.modeling_albert import (
      ALBERT_INPUTS_DOCSTRING,
      ALBERT_START_DOCSTRING,
      AlbertModel,
      AlbertPreTrainedModel,
      AlbertTransformer,
  )
  
  
  SCREAMING_SNAKE_CASE__				:    Dict         =    logging.getLogger(__name__)
  
  
  class 		_SCREAMING_SNAKE_CASE    (	A    ):
  
  
       def 					__snake_case(     self       ,				A_       ,				A_       ,				A_=None       ,				A_=None						):
         if current_layer == 0:
           _UpperCAmelCase		:       Any       =    self.embedding_hidden_mapping_in(A_						)
         else:
           _UpperCAmelCase		:       Optional[int]       =    hidden_states[0]
  
         _UpperCAmelCase		:       Tuple       =    int(self.config.num_hidden_layers / self.config.num_hidden_groups						)
  
         # Index of the hidden group
         _UpperCAmelCase		:       Union[str, Any]       =    int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)						)
  
         _UpperCAmelCase		:       int       =    self.albert_layer_groups[group_idx](
             A_       ,				A_       ,				head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]       ,				)
         _UpperCAmelCase		:       Dict       =    layer_group_output[0]
  
         return (hidden_states,)
  
  
  @add_start_docstrings(
      '''The bare ALBERT Model transformer with PABEE outputting raw hidden-states without any specific head on top.'''       ,  A       ,  )
  class 		_SCREAMING_SNAKE_CASE    (	A    ):
  
  
       def __init__(     self       ,				A_						):
         super().__init__(A_						)
  
         _UpperCAmelCase		:       List[Any]       =    AlbertTransformerWithPabee(A_						)
  
         self.init_weights()
         _UpperCAmelCase		:       Tuple       =    0
         _UpperCAmelCase		:       Optional[Any]       =    0
         _UpperCAmelCase		:       int       =    0
  
         _UpperCAmelCase		:       Dict       =    0
  
  
       def 					__snake_case(     self       ,				A_						):
         _UpperCAmelCase		:       str       =    threshold
  
  
       def 					__snake_case(     self       ,				A_						):
         _UpperCAmelCase		:       Optional[Any]       =    patience
  
  
       def 					__snake_case(     self						):
         _UpperCAmelCase		:       Tuple       =    0
         _UpperCAmelCase		:       Tuple       =    0
  
  
       def 					__snake_case(     self						):
         _UpperCAmelCase		:       Dict       =    self.inference_layers_num / self.inference_instances_num
         _UpperCAmelCase		:       Optional[int]       =    (
             F'''*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ='''
             F''' {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***'''
         )
         print(A_						)
  
  
  
  
       @add_start_docstrings_to_model_forward(A_						)
       def 					__snake_case(     self       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=False       ,				):
  
         if input_ids is not None and inputs_embeds is not None:
           raise ValueError("""You cannot specify both input_ids and inputs_embeds at the same time"""						)
         elif input_ids is not None:
           _UpperCAmelCase		:       List[str]       =    input_ids.size()
         elif inputs_embeds is not None:
           _UpperCAmelCase		:       Optional[Any]       =    inputs_embeds.size()[:-1]
         else:
           raise ValueError("""You have to specify either input_ids or inputs_embeds"""						)
  
         _UpperCAmelCase		:       List[str]       =    input_ids.device if input_ids is not None else inputs_embeds.device
  
         if attention_mask is None:
           _UpperCAmelCase		:       Any       =    torch.ones(A_       ,				device=A_						)
         if token_type_ids is None:
           _UpperCAmelCase		:       List[str]       =    torch.zeros(A_       ,				dtype=torch.long       ,				device=A_						)
  
         _UpperCAmelCase		:       str       =    attention_mask.unsqueeze(1						).unsqueeze(2						)
         _UpperCAmelCase		:       Tuple       =    extended_attention_mask.to(dtype=self.dtype						)  # fp16 compatibility
         _UpperCAmelCase		:       Tuple       =    (1.0 - extended_attention_mask) * -1_0_0_0_0.0
         _UpperCAmelCase		:       Union[str, Any]       =    self.get_head_mask(A_       ,				self.config.num_hidden_layers						)
  
         _UpperCAmelCase		:       Dict       =    self.embeddings(
             A_       ,				position_ids=A_       ,				token_type_ids=A_       ,				inputs_embeds=A_						)
         _UpperCAmelCase		:       str       =    embedding_output
  
         if self.training:
           _UpperCAmelCase		:       Any       =    []
           for i in range(self.config.num_hidden_layers						):
             _UpperCAmelCase		:       Tuple       =    self.encoder.adaptive_forward(
                 A_       ,				current_layer=A_       ,				attention_mask=A_       ,				head_mask=A_       ,				)
  
             _UpperCAmelCase		:       List[Any]       =    self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]						)						)
             _UpperCAmelCase		:       List[str]       =    output_layers[i](output_dropout(A_						)						)
             res.append(A_						)
         elif self.patience == 0:  # Use all layers for inference
           _UpperCAmelCase		:       Optional[int]       =    self.encoder(A_       ,				A_       ,				head_mask=A_						)
           _UpperCAmelCase		:       Optional[int]       =    self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]						)						)
           _UpperCAmelCase		:       List[Any]       =    [output_layers[self.config.num_hidden_layers - 1](A_						)]
         else:
           _UpperCAmelCase		:       str       =    0
           _UpperCAmelCase		:       str       =    None
           _UpperCAmelCase		:       str       =    0
           for i in range(self.config.num_hidden_layers						):
             calculated_layer_num += 1
             _UpperCAmelCase		:       int       =    self.encoder.adaptive_forward(
                 A_       ,				current_layer=A_       ,				attention_mask=A_       ,				head_mask=A_       ,				)
  
             _UpperCAmelCase		:       Tuple       =    self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]						)						)
             _UpperCAmelCase		:       List[str]       =    output_layers[i](A_						)
             if regression:
               _UpperCAmelCase		:       Optional[Any]       =    logits.detach()
               if patient_result is not None:
                 _UpperCAmelCase		:       str       =    patient_result.detach()
               if (patient_result is not None) and torch.abs(patient_result - labels						) < self.regression_threshold:
                 patient_counter += 1
               else:
                 _UpperCAmelCase		:       Union[str, Any]       =    0
             else:
               _UpperCAmelCase		:       Dict       =    logits.detach().argmax(dim=1						)
               if patient_result is not None:
                 _UpperCAmelCase		:       Optional[Any]       =    patient_result.detach().argmax(dim=1						)
               if (patient_result is not None) and torch.all(labels.eq(A_						)						):
                 patient_counter += 1
               else:
                 _UpperCAmelCase		:       Dict       =    0
  
             _UpperCAmelCase		:       int       =    logits
             if patient_counter == self.patience:
               break
           _UpperCAmelCase		:       Tuple       =    [patient_result]
           self.inference_layers_num += calculated_layer_num
           self.inference_instances_num += 1
  
         return res
  
  
  @add_start_docstrings(
      '''Albert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of
      the pooled output) e.g. for GLUE tasks. '''       ,  A       ,  )
  class 		_SCREAMING_SNAKE_CASE    (	A    ):
  
  
       def __init__(     self       ,				A_						):
         super().__init__(A_						)
         _UpperCAmelCase		:       List[str]       =    config.num_labels
  
         _UpperCAmelCase		:       List[Any]       =    AlbertModelWithPabee(A_						)
         _UpperCAmelCase		:       Optional[int]       =    nn.Dropout(config.classifier_dropout_prob						)
         _UpperCAmelCase		:       Tuple       =    nn.ModuleList(
             [nn.Linear(config.hidden_size       ,				self.config.num_labels						) for _ in range(config.num_hidden_layers						)]						)
  
         self.init_weights()
  
  
  
  
       @add_start_docstrings_to_model_forward(A_						)
       def 					__snake_case(     self       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				A_=None       ,				):
         _UpperCAmelCase		:       Optional[Any]       =    self.albert(
             input_ids=A_       ,				attention_mask=A_       ,				token_type_ids=A_       ,				position_ids=A_       ,				head_mask=A_       ,				inputs_embeds=A_       ,				output_dropout=self.dropout       ,				output_layers=self.classifiers       ,				regression=self.num_labels == 1       ,				)
  
         _UpperCAmelCase		:       str       =    (logits[-1],)
  
         if labels is not None:
           _UpperCAmelCase		:       List[Any]       =    None
           _UpperCAmelCase		:       Optional[int]       =    0
           for ix, logits_item in enumerate(A_						):
             if self.num_labels == 1:
               #  We are doing regression
               _UpperCAmelCase		:       int       =    MSELoss()
               _UpperCAmelCase		:       Union[str, Any]       =    loss_fct(logits_item.view(-1						)       ,				labels.view(-1						)						)
             else:
               _UpperCAmelCase		:       List[Any]       =    CrossEntropyLoss()
               _UpperCAmelCase		:       List[Any]       =    loss_fct(logits_item.view(-1       ,				self.num_labels						)       ,				labels.view(-1						)						)
             if total_loss is None:
               _UpperCAmelCase		:       List[str]       =    loss
             else:
               total_loss += loss * (ix + 1)
             total_weights += ix + 1
           _UpperCAmelCase		:       List[str]       =    (total_loss / total_weights,) + outputs
  
         return outputs
  
  
  
  
  
  