
"""simple docstring"""
from dataclasses import dataclass
from typing import Dict, Optional, Union

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .attention_processor import AttentionProcessor, AttnProcessor
from .embeddings import TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
@dataclass
class       lowerCAmelCase_		(				A__ ):



							'''simple docstring'''





							_snake_case						      =		42
class       lowerCAmelCase_		(				A__				,				A__ ):



							'''simple docstring'''





							@register_to_config
							def __init__( self    ,				snake_case_ = 32    ,				snake_case_ = 64    ,				snake_case_ = 20    ,				snake_case_ = 768    ,				snake_case_=77    ,				snake_case_=4    ,				snake_case_ = 0.0    ,				snake_case_ = "silu"    ,				snake_case_ = None    ,				snake_case_ = None    ,				snake_case_ = "linear"    ,				snake_case_ = "prd"    ,				snake_case_ = None    ,				snake_case_ = None    ,				snake_case_ = None    ,				)       ->  Any:
													super().__init__()
													__lowerCAmelCase     =  num_attention_heads
													__lowerCAmelCase     =  attention_head_dim
													__lowerCAmelCase     =  num_attention_heads * attention_head_dim
													__lowerCAmelCase     =  additional_embeddings

													__lowerCAmelCase     =  time_embed_dim or inner_dim
													__lowerCAmelCase     =  embedding_proj_dim or embedding_dim
													__lowerCAmelCase     =  clip_embed_dim or embedding_dim

													__lowerCAmelCase     =  Timesteps(snake_case_    ,				snake_case_    ,				0				)
													__lowerCAmelCase     =  TimestepEmbedding(snake_case_    ,				snake_case_    ,				out_dim=snake_case_    ,				act_fn=snake_case_				)

													__lowerCAmelCase     =  nn.Linear(snake_case_    ,				snake_case_				)

													if embedding_proj_norm_type is None:
																			__lowerCAmelCase     =  None
													elif embedding_proj_norm_type == "layer":
																			__lowerCAmelCase     =  nn.LayerNorm(snake_case_				)
													else:
																			raise ValueError(f"""unsupported embedding_proj_norm_type: {embedding_proj_norm_type}"""				)

													__lowerCAmelCase     =  nn.Linear(snake_case_    ,				snake_case_				)

													if encoder_hid_proj_type is None:
																			__lowerCAmelCase     =  None
													elif encoder_hid_proj_type == "linear":
																			__lowerCAmelCase     =  nn.Linear(snake_case_    ,				snake_case_				)
													else:
																			raise ValueError(f"""unsupported encoder_hid_proj_type: {encoder_hid_proj_type}"""				)

													__lowerCAmelCase     =  nn.Parameter(torch.zeros(1    ,				num_embeddings + additional_embeddings    ,				snake_case_				)				)

													if added_emb_type == "prd":
																			__lowerCAmelCase     =  nn.Parameter(torch.zeros(1    ,				1    ,				snake_case_				)				)
													elif added_emb_type is None:
																			__lowerCAmelCase     =  None
													else:
																			raise ValueError(
																			    f"""`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `'prd'` or `None`."""				)

													__lowerCAmelCase     =  nn.ModuleList(
													    [
													        BasicTransformerBlock(
													            snake_case_    ,				snake_case_    ,				snake_case_    ,				dropout=snake_case_    ,				activation_fn="""gelu"""    ,				attention_bias=snake_case_    ,				)
													        for d in range(snake_case_				)
													    ]				)

													if norm_in_type == "layer":
																			__lowerCAmelCase     =  nn.LayerNorm(snake_case_				)
													elif norm_in_type is None:
																			__lowerCAmelCase     =  None
													else:
																			raise ValueError(f"""Unsupported norm_in_type: {norm_in_type}."""				)

													__lowerCAmelCase     =  nn.LayerNorm(snake_case_				)

													__lowerCAmelCase     =  nn.Linear(snake_case_    ,				snake_case_				)

													__lowerCAmelCase     =  torch.full(
													    [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]    ,				-10_000.0				)
													causal_attention_mask.triu_(1				)
													__lowerCAmelCase     =  causal_attention_mask[None, ...]
													self.register_buffer("""causal_attention_mask"""    ,				snake_case_    ,				persistent=snake_case_				)

													__lowerCAmelCase     =  nn.Parameter(torch.zeros(1    ,				snake_case_				)				)
													__lowerCAmelCase     =  nn.Parameter(torch.zeros(1    ,				snake_case_				)				)




							@property
							# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
							def    A__						( self				)       ->  Dict[str, AttentionProcessor]:
													__lowerCAmelCase     =  {}

													def fn_recursive_add_processors(snake_case_    ,				snake_case_    ,				snake_case_				):
																			if hasattr(snake_case_    ,				"""set_processor"""				):
																									__lowerCAmelCase     =  module.processor

																			for sub_name, child in module.named_children():
																									fn_recursive_add_processors(f"""{name}.{sub_name}"""    ,				snake_case_    ,				snake_case_				)

																			return processors

													for name, module in self.named_children():
																			fn_recursive_add_processors(snake_case_    ,				snake_case_    ,				snake_case_				)

													return processors




							def    A__						( self    ,				snake_case_				)       ->  Dict:
													__lowerCAmelCase     =  len(self.attn_processors.keys()				)

													if isinstance(snake_case_    ,				snake_case_				) and len(snake_case_				) != count:
																			raise ValueError(
																			    f"""A dict of processors was passed, but the number of processors {len(snake_case_				)} does not match the"""
																			    f""" number of attention layers: {count}. Please make sure to pass {count} processor classes."""				)

													def fn_recursive_attn_processor(snake_case_    ,				snake_case_    ,				snake_case_				):
																			if hasattr(snake_case_    ,				"""set_processor"""				):
																									if not isinstance(snake_case_    ,				snake_case_				):
																															module.set_processor(snake_case_				)
																									else:
																															module.set_processor(processor.pop(f"""{name}.processor"""				)				)

																			for sub_name, child in module.named_children():
																									fn_recursive_attn_processor(f"""{name}.{sub_name}"""    ,				snake_case_    ,				snake_case_				)

													for name, module in self.named_children():
																			fn_recursive_attn_processor(snake_case_    ,				snake_case_    ,				snake_case_				)




							def    A__						( self				)       ->  Optional[Any]:
													self.set_attn_processor(AttnProcessor()				)




							def    A__						( self    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_ = None    ,				snake_case_ = None    ,				snake_case_ = True    ,				)       ->  List[str]:
													__lowerCAmelCase     =  hidden_states.shape[0]

													__lowerCAmelCase     =  timestep
													if not torch.is_tensor(snake_case_				):
																			__lowerCAmelCase     =  torch.tensor([timesteps]    ,				dtype=torch.long    ,				device=hidden_states.device				)
													elif torch.is_tensor(snake_case_				) and len(timesteps.shape				) == 0:
																			__lowerCAmelCase     =  timesteps[None].to(hidden_states.device				)

													# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
													__lowerCAmelCase     =  timesteps * torch.ones(snake_case_    ,				dtype=timesteps.dtype    ,				device=timesteps.device				)

													__lowerCAmelCase     =  self.time_proj(snake_case_				)

													# timesteps does not contain any weights and will always return f32 tensors
													# but time_embedding might be fp16, so we need to cast here.
													__lowerCAmelCase     =  timesteps_projected.to(dtype=self.dtype				)
													__lowerCAmelCase     =  self.time_embedding(snake_case_				)

													if self.embedding_proj_norm is not None:
																			__lowerCAmelCase     =  self.embedding_proj_norm(snake_case_				)

													__lowerCAmelCase     =  self.embedding_proj(snake_case_				)
													if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
																			__lowerCAmelCase     =  self.encoder_hidden_states_proj(snake_case_				)
													elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None:
																			raise ValueError("""`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set"""				)

													__lowerCAmelCase     =  self.proj_in(snake_case_				)

													__lowerCAmelCase     =  self.positional_embedding.to(hidden_states.dtype				)

													__lowerCAmelCase     =  []
													__lowerCAmelCase     =  0

													if encoder_hidden_states is not None:
																			additional_embeds.append(snake_case_				)
																			additional_embeddings_len += encoder_hidden_states.shape[1]

													if len(proj_embeddings.shape				) == 2:
																			__lowerCAmelCase     =  proj_embeddings[:, None, :]

													if len(hidden_states.shape				) == 2:
																			__lowerCAmelCase     =  hidden_states[:, None, :]

													__lowerCAmelCase     =  additional_embeds + [
													    proj_embeddings,
													    time_embeddings[:, None, :],
													    hidden_states,
													]

													if self.prd_embedding is not None:
																			__lowerCAmelCase     =  self.prd_embedding.to(hidden_states.dtype				).expand(snake_case_    ,				-1    ,				-1				)
																			additional_embeds.append(snake_case_				)

													__lowerCAmelCase     =  torch.cat(
													    snake_case_    ,				dim=1    ,				)

													# Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
													__lowerCAmelCase     =  additional_embeddings_len + proj_embeddings.shape[1] + 1
													if positional_embeddings.shape[1] < hidden_states.shape[1]:
																			__lowerCAmelCase     =  F.pad(
																			    snake_case_    ,				(
																			        0,
																			        0,
																			        additional_embeddings_len,
																			        self.prd_embedding.shape[1] if self.prd_embedding is not None else 0,
																			    )    ,				value=0.0    ,				)

													__lowerCAmelCase     =  hidden_states + positional_embeddings

													if attention_mask is not None:
																			__lowerCAmelCase     =  (1 - attention_mask.to(hidden_states.dtype				)) * -10_000.0
																			__lowerCAmelCase     =  F.pad(snake_case_    ,				(0, self.additional_embeddings)    ,				value=0.0				)
																			__lowerCAmelCase     =  (attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype				)
																			__lowerCAmelCase     =  attention_mask.repeat_interleave(self.config.num_attention_heads    ,				dim=0				)

													if self.norm_in is not None:
																			__lowerCAmelCase     =  self.norm_in(snake_case_				)

													for block in self.transformer_blocks:
																			__lowerCAmelCase     =  block(snake_case_    ,				attention_mask=snake_case_				)

													__lowerCAmelCase     =  self.norm_out(snake_case_				)

													if self.prd_embedding is not None:
																			__lowerCAmelCase     =  hidden_states[:, -1]
													else:
																			__lowerCAmelCase     =  hidden_states[:, additional_embeddings_len:]

													__lowerCAmelCase     =  self.proj_to_clip_embeddings(snake_case_				)

													if not return_dict:
																			return (predicted_image_embedding,)

													return PriorTransformerOutput(predicted_image_embedding=snake_case_				)




							def    A__						( self    ,				snake_case_				)       ->  Tuple:
													__lowerCAmelCase     =  (prior_latents * self.clip_std) + self.clip_mean
													return prior_latents


"""simple docstring"""
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.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEmbeddings,
    BertLayer,
    BertPooler,
    BertPreTrainedModel,
)

def  lowercase     (_lowerCAmelCase						):
						__lowerCAmelCase     =  torch.exp(_lowerCAmelCase						)
						__lowerCAmelCase     =  torch.sum(_lowerCAmelCase ,    dim=1						)  # sum of exp(x_i)
						__lowerCAmelCase     =  torch.sum(x * exp_x ,    dim=1						)  # sum of x_i * exp(x_i)
						return torch.log(_lowerCAmelCase						) - B / A
class       lowerCAmelCase_		(				nn.Module ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_				)       ->  Optional[Any]:
													super().__init__()
													__lowerCAmelCase     =  config.output_attentions
													__lowerCAmelCase     =  config.output_hidden_states
													__lowerCAmelCase     =  nn.ModuleList([BertLayer(snake_case_				) for _ in range(config.num_hidden_layers				)]				)
													__lowerCAmelCase     =  nn.ModuleList([BertHighway(snake_case_				) for _ in range(config.num_hidden_layers				)]				)

													__lowerCAmelCase     =  [-1 for _ in range(config.num_hidden_layers				)]




							def    A__						( self    ,				snake_case_				)       ->  List[Any]:
													if (type(snake_case_				) is float) or (type(snake_case_				) is int):
																			for i in range(len(self.early_exit_entropy				)				):
																									__lowerCAmelCase     =  x
													else:
																			__lowerCAmelCase     =  x




							def    A__						( self    ,				snake_case_				)       ->  Dict:
													__lowerCAmelCase     =  pooler.state_dict()
													for highway in self.highway:
																			for name, param in highway.pooler.state_dict().items():
																									param.copy_(loaded_model[name]				)




							def    A__						( self    ,				snake_case_    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				)       ->  List[str]:
													__lowerCAmelCase     =  ()
													__lowerCAmelCase     =  ()
													__lowerCAmelCase     =  ()
													for i, layer_module in enumerate(self.layer				):
																			if self.output_hidden_states:
																									__lowerCAmelCase     =  all_hidden_states + (hidden_states,)

																			__lowerCAmelCase     =  layer_module(
																			    snake_case_    ,				snake_case_    ,				head_mask[i]    ,				snake_case_    ,				snake_case_				)
																			__lowerCAmelCase     =  layer_outputs[0]

																			if self.output_attentions:
																									__lowerCAmelCase     =  all_attentions + (layer_outputs[1],)

																			__lowerCAmelCase     =  (hidden_states,)
																			if self.output_hidden_states:
																									__lowerCAmelCase     =  current_outputs + (all_hidden_states,)
																			if self.output_attentions:
																									__lowerCAmelCase     =  current_outputs + (all_attentions,)

																			__lowerCAmelCase     =  self.highway[i](snake_case_				)
																			# logits, pooled_output

																			if not self.training:
																									__lowerCAmelCase     =  highway_exit[0]
																									__lowerCAmelCase     =  entropy(snake_case_				)
																									__lowerCAmelCase     =  highway_exit + (highway_entropy,)  # logits, hidden_states(?), entropy
																									__lowerCAmelCase     =  all_highway_exits + (highway_exit,)

																									if highway_entropy < self.early_exit_entropy[i]:
																															__lowerCAmelCase     =  (highway_logits,) + current_outputs[1:] + (all_highway_exits,)
																															raise HighwayException(snake_case_    ,				i + 1				)
																			else:
																									__lowerCAmelCase     =  all_highway_exits + (highway_exit,)

        # Add last layer
													if self.output_hidden_states:
																			__lowerCAmelCase     =  all_hidden_states + (hidden_states,)

													__lowerCAmelCase     =  (hidden_states,)
													if self.output_hidden_states:
																			__lowerCAmelCase     =  outputs + (all_hidden_states,)
													if self.output_attentions:
																			__lowerCAmelCase     =  outputs + (all_attentions,)

													__lowerCAmelCase     =  outputs + (all_highway_exits,)
													return outputs  # last-layer hidden state, (all hidden states), (all attentions), all highway exits
@add_start_docstrings(
    '''The Bert Model transformer with early exiting (DeeBERT). '''				,				A__				,				)
class       lowerCAmelCase_		(				A__ ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_				)       ->  Tuple:
													super().__init__(snake_case_				)
													__lowerCAmelCase     =  config

													__lowerCAmelCase     =  BertEmbeddings(snake_case_				)
													__lowerCAmelCase     =  DeeBertEncoder(snake_case_				)
													__lowerCAmelCase     =  BertPooler(snake_case_				)

													self.init_weights()




							def    A__						( self				)       ->  Dict:
													self.encoder.init_highway_pooler(self.pooler				)




							def    A__						( self				)       ->  List[Any]:
													return self.embeddings.word_embeddings




							def    A__						( self    ,				snake_case_				)       ->  Tuple:
													__lowerCAmelCase     =  value




							def    A__						( self    ,				snake_case_				)       ->  List[str]:
													for layer, heads in heads_to_prune.items():
																			self.encoder.layer[layer].attention.prune_heads(snake_case_				)




							@add_start_docstrings_to_model_forward(snake_case_				)
							def    A__						( self    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				)       ->  Optional[Any]:
													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     =  input_ids.size()
													elif inputs_embeds is not None:
																			__lowerCAmelCase     =  inputs_embeds.size()[:-1]
													else:
																			raise ValueError("""You have to specify either input_ids or inputs_embeds"""				)

													__lowerCAmelCase     =  input_ids.device if input_ids is not None else inputs_embeds.device

													if attention_mask is None:
																			__lowerCAmelCase     =  torch.ones(snake_case_    ,				device=snake_case_				)
													if encoder_attention_mask is None:
																			__lowerCAmelCase     =  torch.ones(snake_case_    ,				device=snake_case_				)
													if token_type_ids is None:
																			__lowerCAmelCase     =  torch.zeros(snake_case_    ,				dtype=torch.long    ,				device=snake_case_				)

													# We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
													# ourselves in which case we just need to make it broadcastable to all heads.
													__lowerCAmelCase     =  self.get_extended_attention_mask(snake_case_    ,				snake_case_    ,				snake_case_				)

													# If a 2D ou 3D attention mask is provided for the cross-attention
													# we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
													if encoder_attention_mask.dim() == 3:
																			__lowerCAmelCase     =  encoder_attention_mask[:, None, :, :]
													if encoder_attention_mask.dim() == 2:
																			__lowerCAmelCase     =  encoder_attention_mask[:, None, None, :]

													__lowerCAmelCase     =  encoder_extended_attention_mask.to(
													    dtype=next(self.parameters()				).dtype				)  # fp16 compatibility
													__lowerCAmelCase     =  (1.0 - encoder_extended_attention_mask) * -10_000.0

													# Prepare head mask if needed
													# 1.0 in head_mask indicate we keep the head
													# attention_probs has shape bsz x n_heads x N x N
													# input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
													# and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
													__lowerCAmelCase     =  self.get_head_mask(snake_case_    ,				self.config.num_hidden_layers				)

													__lowerCAmelCase     =  self.embeddings(
													    input_ids=snake_case_    ,				position_ids=snake_case_    ,				token_type_ids=snake_case_    ,				inputs_embeds=snake_case_				)
													__lowerCAmelCase     =  self.encoder(
													    snake_case_    ,				attention_mask=snake_case_    ,				head_mask=snake_case_    ,				encoder_hidden_states=snake_case_    ,				encoder_attention_mask=snake_case_    ,				)
													__lowerCAmelCase     =  encoder_outputs[0]
													__lowerCAmelCase     =  self.pooler(snake_case_				)

													__lowerCAmelCase     =  (
													    sequence_output,
													    pooled_output,
													) + encoder_outputs[
													    1:
													]  # add hidden_states and attentions if they are here
													return outputs  # sequence_output, pooled_output, (hidden_states), (attentions), highway exits
class       lowerCAmelCase_		(				A__ ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_    ,				snake_case_				)       ->  Optional[int]:
													__lowerCAmelCase     =  message
													__lowerCAmelCase     =  exit_layer  # start from 1!
class       lowerCAmelCase_		(				nn.Module ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_				)       ->  List[Any]:
													super().__init__()
													__lowerCAmelCase     =  BertPooler(snake_case_				)
													__lowerCAmelCase     =  nn.Dropout(config.hidden_dropout_prob				)
													__lowerCAmelCase     =  nn.Linear(config.hidden_size    ,				config.num_labels				)




							def    A__						( self    ,				snake_case_				)       ->  Optional[int]:
													# Pooler
													__lowerCAmelCase     =  encoder_outputs[0]
													__lowerCAmelCase     =  self.pooler(snake_case_				)
													# "return" pooler_output

													# BertModel
													__lowerCAmelCase     =  (pooler_input, pooler_output) + encoder_outputs[1:]
													# "return" bmodel_output

													# Dropout and classification
													__lowerCAmelCase     =  bmodel_output[1]

													__lowerCAmelCase     =  self.dropout(snake_case_				)
													__lowerCAmelCase     =  self.classifier(snake_case_				)

													return logits, pooled_output
@add_start_docstrings(
    '''Bert Model (with early exiting - DeeBERT) with a classifier on top,
    also takes care of multi-layer training. '''				,				A__				,				)
class       lowerCAmelCase_		(				A__ ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_				)       ->  Optional[Any]:
													super().__init__(snake_case_				)
													__lowerCAmelCase     =  config.num_labels
													__lowerCAmelCase     =  config.num_hidden_layers

													__lowerCAmelCase     =  DeeBertModel(snake_case_				)
													__lowerCAmelCase     =  nn.Dropout(config.hidden_dropout_prob				)
													__lowerCAmelCase     =  nn.Linear(config.hidden_size    ,				self.config.num_labels				)

													self.init_weights()




							@add_start_docstrings_to_model_forward(snake_case_				)
							def    A__						( self    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=-1    ,				snake_case_=False    ,				)       ->  Optional[int]:
													__lowerCAmelCase     =  self.num_layers
													try:
																			__lowerCAmelCase     =  self.bert(
																			    snake_case_    ,				attention_mask=snake_case_    ,				token_type_ids=snake_case_    ,				position_ids=snake_case_    ,				head_mask=snake_case_    ,				inputs_embeds=snake_case_    ,				)
																			# sequence_output, pooled_output, (hidden_states), (attentions), highway exits

																			__lowerCAmelCase     =  outputs[1]

																			__lowerCAmelCase     =  self.dropout(snake_case_				)
																			__lowerCAmelCase     =  self.classifier(snake_case_				)
																			__lowerCAmelCase     =  (logits,) + outputs[2:]  # add hidden states and attention if they are here
													except HighwayException as e:
																			__lowerCAmelCase     =  e.message
																			__lowerCAmelCase     =  e.exit_layer
																			__lowerCAmelCase     =  outputs[0]

													if not self.training:
																			__lowerCAmelCase     =  entropy(snake_case_				)
																			__lowerCAmelCase     =  []
																			__lowerCAmelCase     =  []
													if labels is not None:
																			if self.num_labels == 1:
																									#  We are doing regression
																									__lowerCAmelCase     =  MSELoss()
																									__lowerCAmelCase     =  loss_fct(logits.view(-1				)    ,				labels.view(-1				)				)
																			else:
																									__lowerCAmelCase     =  CrossEntropyLoss()
																									__lowerCAmelCase     =  loss_fct(logits.view(-1    ,				self.num_labels				)    ,				labels.view(-1				)				)

																			# work with highway exits
																			__lowerCAmelCase     =  []
																			for highway_exit in outputs[-1]:
																									__lowerCAmelCase     =  highway_exit[0]
																									if not self.training:
																															highway_logits_all.append(snake_case_				)
																															highway_entropy.append(highway_exit[2]				)
																									if self.num_labels == 1:
																															#  We are doing regression
																															__lowerCAmelCase     =  MSELoss()
																															__lowerCAmelCase     =  loss_fct(highway_logits.view(-1				)    ,				labels.view(-1				)				)
																									else:
																															__lowerCAmelCase     =  CrossEntropyLoss()
																															__lowerCAmelCase     =  loss_fct(highway_logits.view(-1    ,				self.num_labels				)    ,				labels.view(-1				)				)
																									highway_losses.append(snake_case_				)

																			if train_highway:
																									__lowerCAmelCase     =  (sum(highway_losses[:-1]				),) + outputs
																									# exclude the final highway, of course
																			else:
																									__lowerCAmelCase     =  (loss,) + outputs
													if not self.training:
																			__lowerCAmelCase     =  outputs + ((original_entropy, highway_entropy), exit_layer)
																			if output_layer >= 0:
																									__lowerCAmelCase     =  (
																									    (outputs[0],) + (highway_logits_all[output_layer],) + outputs[2:]
																									)  # use the highway of the last layer

													return outputs  # (loss), logits, (hidden_states), (attentions), (highway_exits)
