code: 
  
  
  
  import torch
  from torch import nn
  from torch.nn import CrossEntropyLoss, MSELoss
  
  from ....modeling_outputs import BaseModelOutputWithPooling, SequenceClassifierOutput
  from ....modeling_utils import ModuleUtilsMixin
  from ....utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
  
  
  _lowercase		:     str 					=logging.get_logger(__name__)
  
  _lowercase		:     Optional[Any] 					="""MMBTConfig"""
  
  
  
  
  
  
  class 	UpperCamelCase_     (     nn.Module       ):
  
  
  
  					def __init__(		self    :    str					,			lowerCamelCase    :    str					,			lowerCamelCase    :    str					,			lowerCamelCase    :    int	):
  							super().__init__()
  							lowerCamelCase_  :    Tuple   	=			config
  							lowerCamelCase_  :    Dict   	=			encoder
  							lowerCamelCase_  :    int   	=			nn.Linear(config.modal_hidden_size					,			config.hidden_size	)
  							lowerCamelCase_  :    str   	=			embeddings.position_embeddings
  							lowerCamelCase_  :    Union[str, Any]   	=			embeddings.token_type_embeddings
  							lowerCamelCase_  :    Optional[int]   	=			embeddings.word_embeddings
  							lowerCamelCase_  :    Any   	=			embeddings.LayerNorm
  							lowerCamelCase_  :    Optional[int]   	=			nn.Dropout(p=config.hidden_dropout_prob	)
  
  
  
  
  
  					def      __a    (		self    :    Tuple					,			lowerCamelCase    :    int					,			lowerCamelCase    :    Any=None					,			lowerCamelCase    :    Optional[Any]=None					,			lowerCamelCase    :    str=None					,			lowerCamelCase    :    int=None	):
  							lowerCamelCase_  :    str   	=			self.proj_embeddings(self.encoder(lowerCamelCase	)	)
  							lowerCamelCase_  :    Optional[int]   	=			token_embeddings.size(1	)
  
  							if start_token is not None:
  									lowerCamelCase_  :    List[Any]   	=			self.word_embeddings(lowerCamelCase	)
  									seq_length += 1
  									lowerCamelCase_  :    Union[str, Any]   	=			torch.cat([start_token_embeds.unsqueeze(1	), token_embeddings]					,			dim=1	)
  
  							if end_token is not None:
  									lowerCamelCase_  :    Any   	=			self.word_embeddings(lowerCamelCase	)
  									seq_length += 1
  									lowerCamelCase_  :    str   	=			torch.cat([token_embeddings, end_token_embeds.unsqueeze(1	)]					,			dim=1	)
  
  							if position_ids is None:
  									lowerCamelCase_  :    Any   	=			torch.arange(lowerCamelCase					,			dtype=torch.long					,			device=input_modal.device	)
  									lowerCamelCase_  :    Optional[Any]   	=			position_ids.unsqueeze(0	).expand(input_modal.size(0	)					,			lowerCamelCase	)
  
  							if token_type_ids is None:
  									lowerCamelCase_  :    int   	=			torch.zeros(
  									    (input_modal.size(0	), seq_length)					,			dtype=torch.long					,			device=input_modal.device	)
  
  							lowerCamelCase_  :    str   	=			self.position_embeddings(lowerCamelCase	)
  							lowerCamelCase_  :    Optional[int]   	=			self.token_type_embeddings(lowerCamelCase	)
  							lowerCamelCase_  :    Dict   	=			token_embeddings + position_embeddings + token_type_embeddings
  							lowerCamelCase_  :    List[str]   	=			self.LayerNorm(lowerCamelCase	)
  							lowerCamelCase_  :    Optional[int]   	=			self.dropout(lowerCamelCase	)
  							return embeddings
  _lowercase		:     Union[str, Any] 					=r"""
      MMBT model was proposed in [Supervised Multimodal Bitransformers for Classifying Images and
      Text](https://github.com/facebookresearch/mmbt) by Douwe Kiela, Suvrat Bhooshan, Hamed Firooz, Davide Testuggine.
      It's a supervised multimodal bitransformer model that fuses information from text and other image encoders, and
      obtain state-of-the-art performance on various multimodal classification benchmark tasks.
  
      This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
      library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
      etc.)
  
      This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
      Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
      and behavior.
  
      Parameters:
          config ([`MMBTConfig`]): Model configuration class with all the parameters of the model.
              Initializing with a config file does not load the weights associated with the model, only the
              configuration.
          transformer (`nn.Module`): A text transformer that is used by MMBT.
              It should have embeddings, encoder, and pooler attributes.
          encoder (`nn.Module`): Encoder for the second modality.
              It should take in a batch of modal inputs and return k, n dimension embeddings.
  """
  
  _lowercase		:     Tuple 					=r"""
      Args:
          input_modal (`torch.FloatTensor` of shape `(batch_size, ***)`):
              The other modality data. It will be the shape that the encoder for that type expects. e.g. With an Image
              Encoder, the shape would be (batch_size, channels, height, width)
          input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
              Indices of input sequence tokens in the vocabulary. It does not expect [CLS] token to be added as it's
              appended to the end of other modality embeddings. Indices can be obtained using [`AutoTokenizer`]. See
              [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details.
  
              [What are input IDs?](../glossary#input-ids)
          modal_start_tokens (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
              Optional start token to be added to Other Modality Embedding. [CLS] Most commonly used for classification
              tasks.
          modal_end_tokens (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
              Optional end token to be added to Other Modality Embedding. [SEP] Most commonly used.
          attention_mask (*optional*) `torch.FloatTensor` of shape `(batch_size, sequence_length)`:
              Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
  
              - 1 for tokens that are **not masked**,
              - 0 for tokens that are **masked**.
  
              [What are attention masks?](../glossary#attention-mask)
          token_type_ids (*optional*) `torch.LongTensor` of shape `(batch_size, sequence_length)`:
              Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
              1]`:
  
              - 0 corresponds to a *sentence A* token,
              - 1 corresponds to a *sentence B* token.
  
              [What are token type IDs?](../glossary#token-type-ids)
          modal_token_type_ids (*optional*) `torch.LongTensor` of shape `(batch_size, modal_sequence_length)`:
              Segment token indices to indicate different portions of the non-text modality. The embeddings from these
              tokens will be summed with the respective token embeddings for the non-text modality.
          position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
              Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
              config.max_position_embeddings - 1]`.
  
              [What are position IDs?](../glossary#position-ids)
          modal_position_ids (`torch.LongTensor` of shape `(batch_size, modal_sequence_length)`, *optional*):
              Indices of positions of each input sequence tokens in the position embeddings for the non-text modality.
              Selected in the range `[0, config.max_position_embeddings - 1]`.
  
              [What are position IDs?](../glossary#position-ids)
          head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
              Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
  
              - 1 indicates the head is **not masked**,
              - 0 indicates the head is **masked**.
  
          inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, embedding_dim)`, *optional*):
              Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
              is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
              model's internal embedding lookup matrix.
          encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
              Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
              the model is configured as a decoder.
          encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
              Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
              the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
  
              - 1 for tokens that are **not masked**,
              - 0 for tokens that are **masked**.
  
          output_attentions (`bool`, *optional*):
              Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
              tensors for more detail.
          output_hidden_states (`bool`, *optional*):
              Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
              more detail.
          return_dict (`bool`, *optional*):
              Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
  """
  
  
  
  
  
  
  @add_start_docstrings(
      'The bare MMBT Model outputting raw hidden-states without any specific head on top.'     ,    snake_case__     ,    )
  class 	UpperCamelCase_     (     nn.Module     ,    snake_case__       ):
  
  
  
  					def __init__(		self    :    str					,			lowerCamelCase    :    Optional[int]					,			lowerCamelCase    :    Optional[Any]					,			lowerCamelCase    :    List[str]	):
  							super().__init__()
  							lowerCamelCase_  :    List[str]   	=			config
  							lowerCamelCase_  :    int   	=			transformer
  							lowerCamelCase_  :    List[Any]   	=			ModalEmbeddings(lowerCamelCase					,			lowerCamelCase					,			transformer.embeddings	)
  
  
  
  					@add_start_docstrings_to_model_forward(lowerCamelCase	)
  					@replace_return_docstrings(output_type=lowerCamelCase					,			config_class=_CONFIG_FOR_DOC	)
  					def      __a    (		self    :    Optional[Any]					,			lowerCamelCase    :    Any					,			lowerCamelCase    :    Union[str, Any]=None					,			lowerCamelCase    :    int=None					,			lowerCamelCase    :    Dict=None					,			lowerCamelCase    :    Optional[Any]=None					,			lowerCamelCase    :    List[str]=None					,			lowerCamelCase    :    Any=None					,			lowerCamelCase    :    str=None					,			lowerCamelCase    :    List[str]=None					,			lowerCamelCase    :    Tuple=None					,			lowerCamelCase    :    Union[str, Any]=None					,			lowerCamelCase    :    Union[str, Any]=None					,			lowerCamelCase    :    Dict=None					,			lowerCamelCase    :    List[Any]=None					,			lowerCamelCase    :    int=None					,			lowerCamelCase    :    int=None					,			):
  							lowerCamelCase_  :    Tuple   	=			output_attentions if output_attentions is not None else self.config.output_attentions
  							lowerCamelCase_  :    Tuple   	=			(
  							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
  							)
  							lowerCamelCase_  :    List[Any]   	=			return_dict if return_dict is not None else self.config.use_return_dict
  
  							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:
  									lowerCamelCase_  :    int   	=			input_ids.size()
  							elif inputs_embeds is not None:
  									lowerCamelCase_  :    str   	=			inputs_embeds.size()[:-1]
  							else:
  									raise ValueError('You have to specify either input_ids or inputs_embeds'	)
  
  							lowerCamelCase_  :    Any   	=			input_ids.device if input_ids is not None else inputs_embeds.device
  
  							lowerCamelCase_  :    int   	=			self.modal_encoder(
  							    lowerCamelCase					,			start_token=lowerCamelCase					,			end_token=lowerCamelCase					,			position_ids=lowerCamelCase					,			token_type_ids=lowerCamelCase					,			)
  
  							lowerCamelCase_  :    Tuple   	=			modal_embeddings.size()[:-1]
  
  							if token_type_ids is None:
  									lowerCamelCase_  :    List[str]   	=			torch.ones(lowerCamelCase					,			dtype=torch.long					,			device=lowerCamelCase	)
  
  							lowerCamelCase_  :    str   	=			self.transformer.embeddings(
  							    input_ids=lowerCamelCase					,			position_ids=lowerCamelCase					,			token_type_ids=lowerCamelCase					,			inputs_embeds=lowerCamelCase	)
  
  							lowerCamelCase_  :    Any   	=			torch.cat([modal_embeddings, txt_embeddings]					,			1	)
  
  							lowerCamelCase_  :    int   	=			embedding_output.size()[:-1]
  
  							if attention_mask is None:
  									lowerCamelCase_  :    Dict   	=			torch.ones(lowerCamelCase					,			device=lowerCamelCase	)
  							else:
  									lowerCamelCase_  :    List[str]   	=			torch.cat(
  									    [torch.ones(lowerCamelCase					,			device=lowerCamelCase					,			dtype=torch.long	), attention_mask]					,			dim=1	)
  							if encoder_attention_mask is None:
  									lowerCamelCase_  :    int   	=			torch.ones(lowerCamelCase					,			device=lowerCamelCase	)
  							else:
  									lowerCamelCase_  :    Tuple   	=			torch.cat(
  									    [torch.ones(lowerCamelCase					,			device=lowerCamelCase	), encoder_attention_mask]					,			dim=1	)
  
  							lowerCamelCase_  :    Any   	=			self.get_extended_attention_mask(lowerCamelCase					,			lowerCamelCase	)
  							lowerCamelCase_  :    Any   	=			self.invert_attention_mask(lowerCamelCase	)
  							lowerCamelCase_  :    int   	=			self.get_head_mask(lowerCamelCase					,			self.config.num_hidden_layers	)
  
  							lowerCamelCase_  :    str   	=			self.transformer.encoder(
  							    lowerCamelCase					,			attention_mask=lowerCamelCase					,			head_mask=lowerCamelCase					,			encoder_hidden_states=lowerCamelCase					,			encoder_attention_mask=lowerCamelCase					,			output_attentions=lowerCamelCase					,			output_hidden_states=lowerCamelCase					,			return_dict=lowerCamelCase					,			)
  
  							lowerCamelCase_  :    int   	=			encoder_outputs[0]
  							lowerCamelCase_  :    Optional[int]   	=			self.transformer.pooler(lowerCamelCase	)
  
  							if not return_dict:
  									return (sequence_output, pooled_output) + encoder_outputs[1:]
  
  							return BaseModelOutputWithPooling(
  							    last_hidden_state=lowerCamelCase					,			pooler_output=lowerCamelCase					,			hidden_states=encoder_outputs.hidden_states					,			attentions=encoder_outputs.attentions					,			)
  
  
  
  					def      __a    (		self    :    Dict	):
  							return self.embeddings.word_embeddings
  
  
  
  
  
  					def      __a    (		self    :    Tuple					,			lowerCamelCase    :    str	):
  							lowerCamelCase_  :    Dict   	=			value
  
  
  
  
  
  
  @add_start_docstrings(
      '\n    MMBT Model with a sequence classification/regression head on top (a linear layer on top of the pooled output)\n    '     ,    snake_case__     ,    snake_case__     ,    )
  class 	UpperCamelCase_     (     nn.Module       ):
  
  
  
  					def __init__(		self    :    Any					,			lowerCamelCase    :    Optional[Any]					,			lowerCamelCase    :    List[str]					,			lowerCamelCase    :    Dict	):
  							super().__init__()
  							lowerCamelCase_  :    Dict   	=			config.num_labels
  
  							lowerCamelCase_  :    Dict   	=			MMBTModel(lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase	)
  							lowerCamelCase_  :    List[Any]   	=			nn.Dropout(config.hidden_dropout_prob	)
  							lowerCamelCase_  :    str   	=			nn.Linear(config.hidden_size					,			config.num_labels	)
  
  
  
  
  
  					def      __a    (		self    :    List[str]					,			lowerCamelCase    :    Tuple					,			lowerCamelCase    :    List[str]=None					,			lowerCamelCase    :    List[Any]=None					,			lowerCamelCase    :    Dict=None					,			lowerCamelCase    :    List[Any]=None					,			lowerCamelCase    :    List[str]=None					,			lowerCamelCase    :    Tuple=None					,			lowerCamelCase    :    Dict=None					,			lowerCamelCase    :    Union[str, Any]=None					,			lowerCamelCase    :    Dict=None					,			lowerCamelCase    :    Tuple=None					,			lowerCamelCase    :    Optional[Any]=None					,			lowerCamelCase    :    Dict=None					,			):
  							lowerCamelCase_  :    Any   	=			return_dict if return_dict is not None else self.config.use_return_dict
  
  							lowerCamelCase_  :    int   	=			self.mmbt(
  							    input_modal=lowerCamelCase					,			input_ids=lowerCamelCase					,			modal_start_tokens=lowerCamelCase					,			modal_end_tokens=lowerCamelCase					,			attention_mask=lowerCamelCase					,			token_type_ids=lowerCamelCase					,			modal_token_type_ids=lowerCamelCase					,			position_ids=lowerCamelCase					,			modal_position_ids=lowerCamelCase					,			head_mask=lowerCamelCase					,			inputs_embeds=lowerCamelCase					,			return_dict=lowerCamelCase					,			)
  
  							lowerCamelCase_  :    Dict   	=			outputs[1]
  
  							lowerCamelCase_  :    Optional[Any]   	=			self.dropout(lowerCamelCase	)
  							lowerCamelCase_  :    Union[str, Any]   	=			self.classifier(lowerCamelCase	)
  
  							lowerCamelCase_  :    Dict   	=			None
  							if labels is not None:
  									if self.num_labels == 1:
  											#  We are doing regression
  											lowerCamelCase_  :    int   	=			MSELoss()
  											lowerCamelCase_  :    str   	=			loss_fct(logits.view(-1	)					,			labels.view(-1	)	)
  									else:
  											lowerCamelCase_  :    Union[str, Any]   	=			CrossEntropyLoss()
  											lowerCamelCase_  :    List[str]   	=			loss_fct(logits.view(-1					,			self.num_labels	)					,			labels.view(-1	)	)
  
  							if not return_dict:
  									lowerCamelCase_  :    Any   	=			(logits,) + outputs[2:]
  									return ((loss,) + output) if loss is not None else output
  
  							return SequenceClassifierOutput(
  							    loss=lowerCamelCase					,			logits=lowerCamelCase					,			hidden_states=outputs.hidden_states					,			attentions=outputs.attentions					,			)
  
style_context: 
  
  
  
  
  
  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
  
  
  _lowercase		:     Optional[Any] 					=logging.get_logger(__name__)
  
  # General docstring
  _lowercase		:     Union[str, Any] 					="""ResNetConfig"""
  
  # Base docstring
  _lowercase		:     Dict 					="""microsoft/resnet-50"""
  _lowercase		:     Tuple 					=[1, 2048, 7, 7]
  
  # Image classification docstring
  _lowercase		:     Dict 					="""microsoft/resnet-50"""
  _lowercase		:     int 					="""tiger cat"""
  
  _lowercase		:     Optional[Any] 					=[
      """microsoft/resnet-50""",
      # See all resnet models at https://huggingface.co/models?filter=resnet
  ]
  
  
  
  
  
  
  class 	UpperCamelCase_     (     tf.keras.layers.Layer       ):
  
  
  
  					def __init__(		self    :    List[str]					,			lowerCamelCase    :    int					,			lowerCamelCase    :    int = 3					,			lowerCamelCase    :    int = 1					,			lowerCamelCase    :    str = "relu"					,			**lowerCamelCase    :    Tuple	):
  							super().__init__(**lowerCamelCase	)
  							lowerCamelCase_  :    Tuple   	=			kernel_size // 2
  							lowerCamelCase_  :    Optional[int]   	=			tf.keras.layers.ConvaD(
  							    lowerCamelCase					,			kernel_size=lowerCamelCase					,			strides=lowerCamelCase					,			padding='valid'					,			use_bias=lowerCamelCase					,			name='convolution'	)
  							# Use same default momentum and epsilon as PyTorch equivalent
  							lowerCamelCase_  :    Union[str, Any]   	=			tf.keras.layers.BatchNormalization(epsilon=1E-5					,			momentum=0.9					,			name='normalization'	)
  							lowerCamelCase_  :    Optional[Any]   	=			ACTaFN[activation] if activation is not None else tf.keras.layers.Activation('linear'	)
  
  
  
  					def      __a    (		self    :    int					,			lowerCamelCase    :    tf.Tensor	):
  							# Pad to match that done in the PyTorch Conv2D model
  							lowerCamelCase_  :    int   	=			(self.pad_value, self.pad_value)
  							lowerCamelCase_  :    Optional[Any]   	=			tf.pad(lowerCamelCase					,			[(0, 0), height_pad, width_pad, (0, 0)]	)
  							lowerCamelCase_  :    str   	=			self.conv(lowerCamelCase	)
  							return hidden_state
  
  
  
  
  
  					def      __a    (		self    :    Tuple					,			lowerCamelCase    :    tf.Tensor					,			lowerCamelCase    :    bool = False	):
  							lowerCamelCase_  :    Optional[Any]   	=			self.convolution(lowerCamelCase	)
  							lowerCamelCase_  :    int   	=			self.normalization(lowerCamelCase					,			training=lowerCamelCase	)
  							lowerCamelCase_  :    Any   	=			self.activation(lowerCamelCase	)
  							return hidden_state
  
  
  
  
  
  
  class 	UpperCamelCase_     (     tf.keras.layers.Layer       ):
  
  
  
  					def __init__(		self    :    Dict					,			lowerCamelCase    :    ResNetConfig					,			**lowerCamelCase    :    Optional[Any]	):
  							super().__init__(**lowerCamelCase	)
  							lowerCamelCase_  :    Union[str, Any]   	=			TFResNetConvLayer(
  							    config.embedding_size					,			kernel_size=7					,			stride=2					,			activation=config.hidden_act					,			name='embedder'					,			)
  							lowerCamelCase_  :    Any   	=			tf.keras.layers.MaxPoolaD(pool_size=3					,			strides=2					,			padding='valid'					,			name='pooler'	)
  							lowerCamelCase_  :    Any   	=			config.num_channels
  
  
  
  
  
  					def      __a    (		self    :    List[Any]					,			lowerCamelCase    :    tf.Tensor					,			lowerCamelCase    :    bool = False	):
  							lowerCamelCase_	,			lowerCamelCase_	,			lowerCamelCase_	,			lowerCamelCase_  :    Tuple   	=			shape_list(lowerCamelCase	)
  							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.'	)
  							lowerCamelCase_  :    Optional[int]   	=			pixel_values
  							lowerCamelCase_  :    Optional[Any]   	=			self.embedder(lowerCamelCase	)
  							lowerCamelCase_  :    Optional[int]   	=			tf.pad(lowerCamelCase					,			[[0, 0], [1, 1], [1, 1], [0, 0]]	)
  							lowerCamelCase_  :    str   	=			self.pooler(lowerCamelCase	)
  							return hidden_state
  
  
  
  
  
  
  class 	UpperCamelCase_     (     tf.keras.layers.Layer       ):
  
  
  
  					def __init__(		self    :    Union[str, Any]					,			lowerCamelCase    :    int					,			lowerCamelCase    :    int = 2					,			**lowerCamelCase    :    List[Any]	):
  							super().__init__(**lowerCamelCase	)
  							lowerCamelCase_  :    int   	=			tf.keras.layers.ConvaD(
  							    lowerCamelCase					,			kernel_size=1					,			strides=lowerCamelCase					,			use_bias=lowerCamelCase					,			name='convolution'	)
  							# Use same default momentum and epsilon as PyTorch equivalent
  							lowerCamelCase_  :    Any   	=			tf.keras.layers.BatchNormalization(epsilon=1E-5					,			momentum=0.9					,			name='normalization'	)
  
  
  
  
  
  					def      __a    (		self    :    Tuple					,			lowerCamelCase    :    tf.Tensor					,			lowerCamelCase    :    bool = False	):
  							lowerCamelCase_  :    Union[str, Any]   	=			x
  							lowerCamelCase_  :    Optional[int]   	=			self.convolution(lowerCamelCase	)
  							lowerCamelCase_  :    List[str]   	=			self.normalization(lowerCamelCase					,			training=lowerCamelCase	)
  							return hidden_state
  
  
  
  
  
  
  class 	UpperCamelCase_     (     tf.keras.layers.Layer       ):
  
  
  
  					def __init__(		self    :    List[Any]					,			lowerCamelCase    :    int					,			lowerCamelCase    :    int					,			lowerCamelCase    :    int = 1					,			lowerCamelCase    :    str = "relu"					,			**lowerCamelCase    :    Optional[int]	):
  							super().__init__(**lowerCamelCase	)
  							lowerCamelCase_  :    Any   	=			in_channels != out_channels or stride != 1
  							lowerCamelCase_  :    Tuple   	=			TFResNetConvLayer(lowerCamelCase					,			stride=lowerCamelCase					,			name='layer.0'	)
  							lowerCamelCase_  :    str   	=			TFResNetConvLayer(lowerCamelCase					,			activation=lowerCamelCase					,			name='layer.1'	)
  							lowerCamelCase_  :    Any   	=			(
  							    TFResNetShortCut(lowerCamelCase					,			stride=lowerCamelCase					,			name='shortcut'	)
  							    if should_apply_shortcut
  							    else tf.keras.layers.Activation('linear'					,			name='shortcut'	)
  							)
  							lowerCamelCase_  :    str   	=			ACTaFN[activation]
  
  
  
  
  
  					def      __a    (		self    :    List[Any]					,			lowerCamelCase    :    tf.Tensor					,			lowerCamelCase    :    bool = False	):
  							lowerCamelCase_  :    Optional[int]   	=			hidden_state
  							lowerCamelCase_  :    Tuple   	=			self.conva(lowerCamelCase					,			training=lowerCamelCase	)
  							lowerCamelCase_  :    List[Any]   	=			self.conva(lowerCamelCase					,			training=lowerCamelCase	)
  							lowerCamelCase_  :    int   	=			self.shortcut(lowerCamelCase					,			training=lowerCamelCase	)
  							hidden_state += residual
  							lowerCamelCase_  :    Dict   	=			self.activation(lowerCamelCase	)
  							return hidden_state
  
  
  
  
  
  
  class 	UpperCamelCase_     (     tf.keras.layers.Layer       ):
  
  
  
  					def __init__(		self    :    Optional[Any]					,			lowerCamelCase    :    int					,			lowerCamelCase    :    int					,			lowerCamelCase    :    int = 1					,			lowerCamelCase    :    str = "relu"					,			lowerCamelCase    :    int = 4					,			**lowerCamelCase    :    Optional[int]					,			):
  							super().__init__(**lowerCamelCase	)
  							lowerCamelCase_  :    List[Any]   	=			in_channels != out_channels or stride != 1
  							lowerCamelCase_  :    Optional[Any]   	=			out_channels // reduction
  							lowerCamelCase_  :    int   	=			TFResNetConvLayer(lowerCamelCase					,			kernel_size=1					,			name='layer.0'	)
  							lowerCamelCase_  :    Tuple   	=			TFResNetConvLayer(lowerCamelCase					,			stride=lowerCamelCase					,			name='layer.1'	)
  							lowerCamelCase_  :    Optional[int]   	=			TFResNetConvLayer(lowerCamelCase					,			kernel_size=1					,			activation=lowerCamelCase					,			name='layer.2'	)
  							lowerCamelCase_  :    Optional[Any]   	=			(
  							    TFResNetShortCut(lowerCamelCase					,			stride=lowerCamelCase					,			name='shortcut'	)
  							    if should_apply_shortcut
  							    else tf.keras.layers.Activation('linear'					,			name='shortcut'	)
  							)
  							lowerCamelCase_  :    List[Any]   	=			ACTaFN[activation]
  
  
  
  
  
  					def      __a    (		self    :    str					,			lowerCamelCase    :    tf.Tensor					,			lowerCamelCase    :    bool = False	):
  							lowerCamelCase_  :    int   	=			hidden_state
  							lowerCamelCase_  :    Tuple   	=			self.conva(lowerCamelCase					,			training=lowerCamelCase	)
  							lowerCamelCase_  :    List[Any]   	=			self.conva(lowerCamelCase					,			training=lowerCamelCase	)
  							lowerCamelCase_  :    List[str]   	=			self.conva(lowerCamelCase					,			training=lowerCamelCase	)
  							lowerCamelCase_  :    List[Any]   	=			self.shortcut(lowerCamelCase					,			training=lowerCamelCase	)
  							hidden_state += residual
  							lowerCamelCase_  :    int   	=			self.activation(lowerCamelCase	)
  							return hidden_state
  
  
  
  
  
  
  class 	UpperCamelCase_     (     tf.keras.layers.Layer       ):
  
  
  
  					def __init__(		self    :    int					,			lowerCamelCase    :    ResNetConfig					,			lowerCamelCase    :    int					,			lowerCamelCase    :    int					,			lowerCamelCase    :    int = 2					,			lowerCamelCase    :    int = 2					,			**lowerCamelCase    :    List[Any]	):
  							super().__init__(**lowerCamelCase	)
  
  							lowerCamelCase_  :    List[Any]   	=			TFResNetBottleNeckLayer if config.layer_type == 'bottleneck' else TFResNetBasicLayer
  
  							lowerCamelCase_  :    Tuple   	=			[layer(lowerCamelCase					,			lowerCamelCase					,			stride=lowerCamelCase					,			activation=config.hidden_act					,			name='layers.0'	)]
  							layers += [
  							    layer(lowerCamelCase					,			lowerCamelCase					,			activation=config.hidden_act					,			name=F"layers.{i + 1}"	)
  							    for i in range(depth - 1	)
  							]
  							lowerCamelCase_  :    Dict   	=			layers
  
  
  
  
  
  					def      __a    (		self    :    Any					,			lowerCamelCase    :    tf.Tensor					,			lowerCamelCase    :    bool = False	):
  							for layer in self.stage_layers:
  									lowerCamelCase_  :    Any   	=			layer(lowerCamelCase					,			training=lowerCamelCase	)
  							return hidden_state
  
  
  
  
  
  
  class 	UpperCamelCase_     (     tf.keras.layers.Layer       ):
  
  
  
  					def __init__(		self    :    List[Any]					,			lowerCamelCase    :    ResNetConfig					,			**lowerCamelCase    :    Optional[Any]	):
  							super().__init__(**lowerCamelCase	)
  							# based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
  							lowerCamelCase_  :    List[Any]   	=			[
  							    TFResNetStage(
  							        lowerCamelCase					,			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(lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase					,			depth=lowerCamelCase					,			name=F"stages.{i + 1}"	)	)
  
  
  
  
  
  					def      __a    (		self    :    Optional[int]					,			lowerCamelCase    :    tf.Tensor					,			lowerCamelCase    :    bool = False					,			lowerCamelCase    :    bool = True					,			lowerCamelCase    :    bool = False					,			):
  							lowerCamelCase_  :    Optional[int]   	=			() if output_hidden_states else None
  
  							for stage_module in self.stages:
  									if output_hidden_states:
  											lowerCamelCase_  :    Any   	=			hidden_states + (hidden_state,)
  
  									lowerCamelCase_  :    Optional[Any]   	=			stage_module(lowerCamelCase					,			training=lowerCamelCase	)
  
  							if output_hidden_states:
  									lowerCamelCase_  :    Union[str, Any]   	=			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=lowerCamelCase					,			hidden_states=lowerCamelCase	)
  
  
  
  
  
  
  class 	UpperCamelCase_     (     snake_case__       ):
  					_a			:	int							=				ResNetConfig
  					_a			:	Optional[Any]							=				'resnet'
  					_a			:	int							=				'pixel_values'
  
  
  
  					@property
  					def      __a    (		self    :    Tuple	):
  							return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 2_24, 2_24)					,			dtype=tf.floataa	)}
  _lowercase		:     Union[str, Any] 					=r"""
      This model is a TensorFlow
      [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a
      regular TensorFlow Module and refer to the TensorFlow documentation for all matter related to general usage and
      behavior.
  
      Parameters:
          config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.
              Initializing with a config file does not load the weights associated with the model, only the
              configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
  """
  
  
  _lowercase		:     Optional[Any] 					=r"""
      Args:
          pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):
              Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
              [`ConvNextImageProcessor.__call__`] for details.
  
          output_hidden_states (`bool`, *optional*):
              Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
              more detail.
          return_dict (`bool`, *optional*):
              Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
  """
  
  
  
  
  
  
  @keras_serializable
  class 	UpperCamelCase_     (     tf.keras.layers.Layer       ):
  					_a			:	int							=				ResNetConfig
  
  
  
  					def __init__(		self    :    Dict					,			lowerCamelCase    :    ResNetConfig					,			**lowerCamelCase    :    List[str]	):
  							super().__init__(**lowerCamelCase	)
  							lowerCamelCase_  :    Tuple   	=			config
  							lowerCamelCase_  :    int   	=			TFResNetEmbeddings(lowerCamelCase					,			name='embedder'	)
  							lowerCamelCase_  :    Union[str, Any]   	=			TFResNetEncoder(lowerCamelCase					,			name='encoder'	)
  							lowerCamelCase_  :    Tuple   	=			tf.keras.layers.GlobalAveragePoolingaD(keepdims=lowerCamelCase	)
  
  
  
  
  
  					@unpack_inputs
  					def      __a    (		self    :    List[str]					,			lowerCamelCase    :    tf.Tensor					,			lowerCamelCase    :    Optional[bool] = None					,			lowerCamelCase    :    Optional[bool] = None					,			lowerCamelCase    :    bool = False					,			):
  							lowerCamelCase_  :    str   	=			(
  							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
  							)
  							lowerCamelCase_  :    Optional[Any]   	=			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)
  							lowerCamelCase_  :    str   	=			tf.transpose(lowerCamelCase					,			perm=[0, 2, 3, 1]	)
  							lowerCamelCase_  :    Any   	=			self.embedder(lowerCamelCase					,			training=lowerCamelCase	)
  
  							lowerCamelCase_  :    str   	=			self.encoder(
  							    lowerCamelCase					,			output_hidden_states=lowerCamelCase					,			return_dict=lowerCamelCase					,			training=lowerCamelCase	)
  
  							lowerCamelCase_  :    Optional[int]   	=			encoder_outputs[0]
  
  							lowerCamelCase_  :    str   	=			self.pooler(lowerCamelCase	)
  
  							# Transpose all the outputs to the NCHW format
  							# (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
  							lowerCamelCase_  :    str   	=			tf.transpose(lowerCamelCase					,			(0, 3, 1, 2)	)
  							lowerCamelCase_  :    Tuple   	=			tf.transpose(lowerCamelCase					,			(0, 3, 1, 2)	)
  							lowerCamelCase_  :    Dict   	=			()
  							for hidden_state in encoder_outputs[1:]:
  									lowerCamelCase_  :    Dict   	=			hidden_states + tuple(tf.transpose(lowerCamelCase					,			(0, 3, 1, 2)	) for h in hidden_state	)
  
  							if not return_dict:
  									return (last_hidden_state, pooled_output) + hidden_states
  
  							lowerCamelCase_  :    Dict   	=			hidden_states if output_hidden_states else None
  
  							return TFBaseModelOutputWithPoolingAndNoAttention(
  							    last_hidden_state=lowerCamelCase					,			pooler_output=lowerCamelCase					,			hidden_states=lowerCamelCase					,			)
  
  
  
  
  
  
  @add_start_docstrings(
      'The bare ResNet model outputting raw features without any specific head on top.'     ,    snake_case__     ,    )
  class 	UpperCamelCase_     (     snake_case__       ):
  
  
  
  					def __init__(		self    :    Any					,			lowerCamelCase    :    ResNetConfig					,			**lowerCamelCase    :    int	):
  							super().__init__(lowerCamelCase					,			**lowerCamelCase	)
  							lowerCamelCase_  :    List[str]   	=			TFResNetMainLayer(config=lowerCamelCase					,			name='resnet'	)
  
  
  
  
  
  					@add_start_docstrings_to_model_forward(lowerCamelCase	)
  					@add_code_sample_docstrings(
  					    checkpoint=_CHECKPOINT_FOR_DOC					,			output_type=lowerCamelCase					,			config_class=_CONFIG_FOR_DOC					,			modality='vision'					,			expected_output=_EXPECTED_OUTPUT_SHAPE					,			)
  					@unpack_inputs
  					def      __a    (		self    :    List[str]					,			lowerCamelCase    :    tf.Tensor					,			lowerCamelCase    :    Optional[bool] = None					,			lowerCamelCase    :    Optional[bool] = None					,			lowerCamelCase    :    bool = False					,			):
  							lowerCamelCase_  :    List[str]   	=			(
  							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
  							)
  							lowerCamelCase_  :    Tuple   	=			return_dict if return_dict is not None else self.config.use_return_dict
  
  							lowerCamelCase_  :    Dict   	=			self.resnet(
  							    pixel_values=lowerCamelCase					,			output_hidden_states=lowerCamelCase					,			return_dict=lowerCamelCase					,			training=lowerCamelCase					,			)
  							return resnet_outputs
  
  
  
  
  
  
  @add_start_docstrings(
      '\n    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    '     ,    snake_case__     ,    )
  class 	UpperCamelCase_     (     snake_case__     ,    snake_case__       ):
  
  
  
  					def __init__(		self    :    Tuple					,			lowerCamelCase    :    ResNetConfig					,			**lowerCamelCase    :    Optional[Any]	):
  							super().__init__(lowerCamelCase					,			**lowerCamelCase	)
  							lowerCamelCase_  :    Tuple   	=			config.num_labels
  							lowerCamelCase_  :    Union[str, Any]   	=			TFResNetMainLayer(lowerCamelCase					,			name='resnet'	)
  							# classification head
  							lowerCamelCase_  :    Dict   	=			(
  							    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      __a    (		self    :    List[str]					,			lowerCamelCase    :    tf.Tensor	):
  							lowerCamelCase_  :    Tuple   	=			tf.keras.layers.Flatten()(lowerCamelCase	)
  							lowerCamelCase_  :    Dict   	=			self.classifier_layer(lowerCamelCase	)
  							return logits
  
  
  
  
  
  					@add_start_docstrings_to_model_forward(lowerCamelCase	)
  					@add_code_sample_docstrings(
  					    checkpoint=_IMAGE_CLASS_CHECKPOINT					,			output_type=lowerCamelCase					,			config_class=_CONFIG_FOR_DOC					,			expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT					,			)
  					@unpack_inputs
  					def      __a    (		self    :    Optional[int]					,			lowerCamelCase    :    tf.Tensor = None					,			lowerCamelCase    :    tf.Tensor = None					,			lowerCamelCase    :    bool = None					,			lowerCamelCase    :    bool = None					,			lowerCamelCase    :    bool = False					,			):
  							lowerCamelCase_  :    List[str]   	=			return_dict if return_dict is not None else self.config.use_return_dict
  
  							lowerCamelCase_  :    Optional[int]   	=			self.resnet(
  							    lowerCamelCase					,			output_hidden_states=lowerCamelCase					,			return_dict=lowerCamelCase					,			training=lowerCamelCase	)
  
  							lowerCamelCase_  :    Union[str, Any]   	=			outputs.pooler_output if return_dict else outputs[1]
  
  							lowerCamelCase_  :    Union[str, Any]   	=			self.classifier(lowerCamelCase	)
  
  							lowerCamelCase_  :    List[Any]   	=			None if labels is None else self.hf_compute_loss(lowerCamelCase					,			lowerCamelCase	)
  
  							if not return_dict:
  									lowerCamelCase_  :    int   	=			(logits,) + outputs[2:]
  									return (loss,) + output if loss is not None else output
  
  							return TFImageClassifierOutputWithNoAttention(loss=lowerCamelCase					,			logits=lowerCamelCase					,			hidden_states=outputs.hidden_states	)
  