'''simple docstring'''

import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTJConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, tooslow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
  import jax
  import jax.numpy as jnp

  from transformers.modeling_flax_pytorch_utils import (
      convert_pytorch_state_dict_to_flax,
      load_flax_weights_in_pytorch_model,
  )
  from transformers.models.gptj.modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel

if is_torch_available():
  import torch


class 					_lowercase						:
      '''simple docstring'''


      def __init__(					self       ,lowerCamelCase_       ,lowerCamelCase_=14       ,lowerCamelCase_=7       ,lowerCamelCase_=True       ,lowerCamelCase_=True       ,lowerCamelCase_=False       ,lowerCamelCase_=True       ,lowerCamelCase_=99       ,lowerCamelCase_=32       ,lowerCamelCase_=4       ,lowerCamelCase_=4       ,lowerCamelCase_=4       ,lowerCamelCase_=37       ,lowerCamelCase_="gelu"       ,lowerCamelCase_=0.1       ,lowerCamelCase_=0.1       ,lowerCamelCase_=512       ,lowerCamelCase_=0.02       ,)      ->      int:



         '''simple docstring'''



         UpperCAmelCase__				:					List[Any] 					=	parent
         UpperCAmelCase__				:					Optional[int] 					=	batch_size
         UpperCAmelCase__				:					Dict 					=	seq_length
         UpperCAmelCase__				:					Dict 					=	is_training
         UpperCAmelCase__				:					Any 					=	use_input_mask
         UpperCAmelCase__				:					Any 					=	use_token_type_ids
         UpperCAmelCase__				:					Any 					=	use_labels
         UpperCAmelCase__				:					Any 					=	vocab_size
         UpperCAmelCase__				:					Optional[Any] 					=	hidden_size
         UpperCAmelCase__				:					Optional[int] 					=	rotary_dim
         UpperCAmelCase__				:					int 					=	num_hidden_layers
         UpperCAmelCase__				:					Any 					=	num_attention_heads
         UpperCAmelCase__				:					Tuple 					=	intermediate_size
         UpperCAmelCase__				:					Dict 					=	hidden_act
         UpperCAmelCase__				:					Tuple 					=	hidden_dropout_prob
         UpperCAmelCase__				:					List[str] 					=	attention_probs_dropout_prob
         UpperCAmelCase__				:					Dict 					=	max_position_embeddings
         UpperCAmelCase__				:					Tuple 					=	initializer_range
         UpperCAmelCase__				:					Union[str, Any] 					=	None
         UpperCAmelCase__				:					List[str] 					=	vocab_size - 1
         UpperCAmelCase__				:					List[Any] 					=	vocab_size - 1
         UpperCAmelCase__				:					int 					=	vocab_size - 1




      def 		lowerCAmelCase__       (					self				)      ->      List[Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[int] 					=	ids_tensor([self.batch_size, self.seq_length]       ,self.vocab_size				)

         UpperCAmelCase__				:					Dict 					=	None
         if self.use_input_mask:
            UpperCAmelCase__				:					int 					=	random_attention_mask([self.batch_size, self.seq_length]				)

         UpperCAmelCase__				:					Optional[int] 					=	GPTJConfig(
             vocab_size=self.vocab_size       ,n_embd=self.hidden_size       ,n_layer=self.num_hidden_layers       ,n_head=self.num_attention_heads       ,n_positions=self.max_position_embeddings       ,use_cache=lowerCamelCase_       ,bos_token_id=self.bos_token_id       ,eos_token_id=self.eos_token_id       ,pad_token_id=self.pad_token_id       ,rotary_dim=self.rotary_dim       ,)

         return (config, input_ids, input_mask)




      def 		lowerCAmelCase__       (					self				)      ->      List[str]:



         '''simple docstring'''



         UpperCAmelCase__				:					str 					=	self.prepare_config_and_inputs()
         UpperCAmelCase__					,   UpperCAmelCase__					,   UpperCAmelCase__				:					List[Any] 					=	config_and_inputs
         UpperCAmelCase__				:					Tuple 					=	{'''input_ids''': input_ids, '''attention_mask''': attention_mask}
         return config, inputs_dict




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      Dict:



         '''simple docstring'''



         UpperCAmelCase__				:					Dict 					=	20
         UpperCAmelCase__				:					List[str] 					=	model_class_name(lowerCamelCase_				)

         UpperCAmelCase__				:					Tuple 					=	model.init_cache(input_ids.shape[0]       ,lowerCamelCase_				)
         UpperCAmelCase__				:					Dict 					=	jnp.ones((input_ids.shape[0], max_decoder_length)       ,dtype='''i4'''				)

         UpperCAmelCase__				:					Tuple 					=	jnp.broadcast_to(
             jnp.arange(input_ids.shape[-1] - 1				)[None, :]       ,(input_ids.shape[0], input_ids.shape[-1] - 1)				)
         UpperCAmelCase__				:					Optional[Any] 					=	model(
             input_ids[:, :-1]       ,attention_mask=lowerCamelCase_       ,past_key_values=lowerCamelCase_       ,position_ids=lowerCamelCase_       ,)

         UpperCAmelCase__				:					int 					=	jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,dtype='''i4'''				)
         UpperCAmelCase__				:					Optional[int] 					=	model(
             input_ids[:, -1:]       ,attention_mask=lowerCamelCase_       ,past_key_values=outputs_cache.past_key_values       ,position_ids=lowerCamelCase_       ,)

         UpperCAmelCase__				:					List[str] 					=	model(lowerCamelCase_				)

         UpperCAmelCase__				:					Dict 					=	np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])				)				)
         self.parent.assertTrue(diff < 1e-3       ,msg=f'''Max diff is {diff}'''				)




      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)      ->      Tuple:



         '''simple docstring'''



         UpperCAmelCase__				:					Union[str, Any] 					=	20
         UpperCAmelCase__				:					Any 					=	model_class_name(lowerCamelCase_				)

         UpperCAmelCase__				:					List[str] 					=	jnp.concatenate(
             [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])				)]       ,axis=-1       ,)

         UpperCAmelCase__				:					Tuple 					=	model.init_cache(input_ids.shape[0]       ,lowerCamelCase_				)
         UpperCAmelCase__				:					List[Any] 					=	jnp.broadcast_to(
             jnp.arange(input_ids.shape[-1] - 1				)[None, :]       ,(input_ids.shape[0], input_ids.shape[-1] - 1)				)

         UpperCAmelCase__				:					Optional[int] 					=	model(
             input_ids[:, :-1]       ,attention_mask=lowerCamelCase_       ,past_key_values=lowerCamelCase_       ,position_ids=lowerCamelCase_       ,)
         UpperCAmelCase__				:					Tuple 					=	jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,dtype='''i4'''				)
         UpperCAmelCase__				:					str 					=	model(
             input_ids[:, -1:]       ,past_key_values=outputs_cache.past_key_values       ,attention_mask=lowerCamelCase_       ,position_ids=lowerCamelCase_       ,)

         UpperCAmelCase__				:					Optional[Any] 					=	model(lowerCamelCase_       ,attention_mask=lowerCamelCase_				)

         UpperCAmelCase__				:					List[Any] 					=	np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])				)				)
         self.parent.assertTrue(diff < 1e-3       ,msg=f'''Max diff is {diff}'''				)




@require_flax
class 					_lowercase						(   lowerCAmelCase		,lowerCAmelCase		,unittest.TestCase ):
      '''simple docstring'''


      UpperCAmelCase_   :   Union[str, Any]   				= (FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
      UpperCAmelCase_   :   Tuple   				= (FlaxGPTJForCausalLM,) if is_flax_available() else ()




      def 		lowerCAmelCase__       (					self				)      ->      Union[str, Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					Union[str, Any] 					=	FlaxGPTJModelTester(self				)




      def 		lowerCAmelCase__       (					self				)      ->      Optional[int]:



         '''simple docstring'''



         for model_class_name in self.all_model_classes:
            UpperCAmelCase__					,   UpperCAmelCase__					,   UpperCAmelCase__				:					List[str] 					=	self.model_tester.prepare_config_and_inputs()
            self.model_tester.check_use_cache_forward(lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)




      def 		lowerCAmelCase__       (					self				)      ->      Any:



         '''simple docstring'''



         for model_class_name in self.all_model_classes:
            UpperCAmelCase__					,   UpperCAmelCase__					,   UpperCAmelCase__				:					Union[str, Any] 					=	self.model_tester.prepare_config_and_inputs()
            self.model_tester.check_use_cache_forward_with_attn_mask(
                lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)




      @tooslow
      def 		lowerCAmelCase__       (					self				)      ->      int:



         '''simple docstring'''



         UpperCAmelCase__				:					List[str] 					=	GPTaTokenizer.from_pretrained('''gpt2'''       ,pad_token='''<|endoftext|>'''       ,padding_side='''left'''				)
         UpperCAmelCase__				:					Union[str, Any] 					=	tokenizer(['''Hello this is a long string''', '''Hey''']       ,return_tensors='''np'''       ,padding=lowerCamelCase_       ,truncation=lowerCamelCase_				)

         UpperCAmelCase__				:					Tuple 					=	FlaxGPTJForCausalLM.from_pretrained('''EleutherAI/gpt-j-6B'''				)
         UpperCAmelCase__				:					int 					=	False
         UpperCAmelCase__				:					Any 					=	model.config.eos_token_id

         UpperCAmelCase__				:					Union[str, Any] 					=	jax.jit(model.generate				)

         UpperCAmelCase__				:					Optional[int] 					=	jit_generate(
             inputs['''input_ids''']       ,attention_mask=inputs['''attention_mask''']       ,pad_token_id=tokenizer.pad_token_id				).sequences

         UpperCAmelCase__				:					Tuple 					=	tokenizer.batch_decode(lowerCamelCase_       ,skip_special_tokens=lowerCamelCase_				)

         UpperCAmelCase__				:					Dict 					=	[
             '''Hello this is a long string of text.\n\nI\'m trying to get the text of the''',
             '''Hey, I\'m a little late to the party. I\'m going to''',
         ]

         self.assertListEqual(lowerCamelCase_       ,lowerCamelCase_				)




      @is_pt_flax_cross_test
      def 		lowerCAmelCase__       (					self				)      ->      Union[str, Any]:



         '''simple docstring'''



         UpperCAmelCase__					,   UpperCAmelCase__				:					List[str] 					=	self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
            with self.subTest(model_class.__name__				):
               # prepare inputs
               UpperCAmelCase__				:					Union[str, Any] 					=	self._prepare_for_class(lowerCamelCase_       ,lowerCamelCase_				)
               UpperCAmelCase__				:					int 					=	{k: torch.tensor(v.tolist()				) for k, v in prepared_inputs_dict.items()}

               # load corresponding PyTorch class
               UpperCAmelCase__				:					Union[str, Any] 					=	model_class.__name__[4:]  # Skip the "Flax" at the beginning
               UpperCAmelCase__				:					Dict 					=	getattr(lowerCamelCase_       ,lowerCamelCase_				)

               UpperCAmelCase__					,   UpperCAmelCase__				:					List[str] 					=	pt_inputs['''input_ids'''].shape
               UpperCAmelCase__				:					str 					=	np.random.randint(0       ,seq_length - 1       ,size=(batch_size,)				)
               for batch_idx, start_index in enumerate(lowerCamelCase_				):
                  UpperCAmelCase__				:					List[Any] 					=	0
                  UpperCAmelCase__				:					int 					=	1
                  UpperCAmelCase__				:					Tuple 					=	0
                  UpperCAmelCase__				:					Optional[Any] 					=	1
               UpperCAmelCase__				:					Optional[Any] 					=	pt_model_class(lowerCamelCase_				).eval()
               UpperCAmelCase__				:					Dict 					=	model_class(lowerCamelCase_       ,dtype=jnp.floataa				)

               UpperCAmelCase__				:					Tuple 					=	convert_pytorch_state_dict_to_flax(pt_model.state_dict()       ,lowerCamelCase_				)
               UpperCAmelCase__				:					str 					=	fx_state

               with torch.no_grad():
                  UpperCAmelCase__				:					Union[str, Any] 					=	pt_model(**lowerCamelCase_				).to_tuple()

               UpperCAmelCase__				:					Optional[Any] 					=	fx_model(**lowerCamelCase_				).to_tuple()
               self.assertEqual(len(lowerCamelCase_				)       ,len(lowerCamelCase_				)       ,'''Output lengths differ between Flax and PyTorch'''				)
               for fx_output, pt_output in zip(lowerCamelCase_       ,lowerCamelCase_				):
                  self.assert_almost_equals(fx_output[:, -1]       ,pt_output[:, -1].numpy()       ,4e-2				)

               with tempfile.TemporaryDirectory() as tmpdirname:
                  pt_model.save_pretrained(lowerCamelCase_				)
                  UpperCAmelCase__				:					Union[str, Any] 					=	model_class.from_pretrained(lowerCamelCase_       ,from_pt=lowerCamelCase_				)

               UpperCAmelCase__				:					Optional[int] 					=	fx_model_loaded(**lowerCamelCase_				).to_tuple()
               self.assertEqual(
                   len(lowerCamelCase_				)       ,len(lowerCamelCase_				)       ,'''Output lengths differ between Flax and PyTorch'''				)
               for fx_output_loaded, pt_output in zip(lowerCamelCase_       ,lowerCamelCase_				):
                  self.assert_almost_equals(fx_output_loaded[:, -1]       ,pt_output[:, -1].numpy()       ,4e-2				)




      @is_pt_flax_cross_test
      def 		lowerCAmelCase__       (					self				)      ->      List[str]:



         '''simple docstring'''



         UpperCAmelCase__					,   UpperCAmelCase__				:					str 					=	self.model_tester.prepare_config_and_inputs_for_common()
         for model_class in self.all_model_classes:
            with self.subTest(model_class.__name__				):
               # prepare inputs
               UpperCAmelCase__				:					Union[str, Any] 					=	self._prepare_for_class(lowerCamelCase_       ,lowerCamelCase_				)
               UpperCAmelCase__				:					Union[str, Any] 					=	{k: torch.tensor(v.tolist()				) for k, v in prepared_inputs_dict.items()}

               # load corresponding PyTorch class
               UpperCAmelCase__				:					Any 					=	model_class.__name__[4:]  # Skip the "Flax" at the beginning
               UpperCAmelCase__				:					int 					=	getattr(lowerCamelCase_       ,lowerCamelCase_				)

               UpperCAmelCase__				:					Optional[Any] 					=	pt_model_class(lowerCamelCase_				).eval()
               UpperCAmelCase__				:					Union[str, Any] 					=	model_class(lowerCamelCase_       ,dtype=jnp.floataa				)

               UpperCAmelCase__				:					Tuple 					=	load_flax_weights_in_pytorch_model(lowerCamelCase_       ,fx_model.params				)
               UpperCAmelCase__					,   UpperCAmelCase__				:					Optional[Any] 					=	pt_inputs['''input_ids'''].shape
               UpperCAmelCase__				:					Optional[Any] 					=	np.random.randint(0       ,seq_length - 1       ,size=(batch_size,)				)
               for batch_idx, start_index in enumerate(lowerCamelCase_				):
                  UpperCAmelCase__				:					Optional[int] 					=	0
                  UpperCAmelCase__				:					Tuple 					=	1
                  UpperCAmelCase__				:					Tuple 					=	0
                  UpperCAmelCase__				:					Optional[Any] 					=	1

               # make sure weights are tied in PyTorch
               pt_model.tie_weights()

               with torch.no_grad():
                  UpperCAmelCase__				:					Union[str, Any] 					=	pt_model(**lowerCamelCase_				).to_tuple()

               UpperCAmelCase__				:					Any 					=	fx_model(**lowerCamelCase_				).to_tuple()
               self.assertEqual(len(lowerCamelCase_				)       ,len(lowerCamelCase_				)       ,'''Output lengths differ between Flax and PyTorch'''				)
               for fx_output, pt_output in zip(lowerCamelCase_       ,lowerCamelCase_				):
                  self.assert_almost_equals(fx_output[:, -1]       ,pt_output[:, -1].numpy()       ,4e-2				)

               with tempfile.TemporaryDirectory() as tmpdirname:
                  fx_model.save_pretrained(lowerCamelCase_				)
                  UpperCAmelCase__				:					Dict 					=	pt_model_class.from_pretrained(lowerCamelCase_       ,from_flax=lowerCamelCase_				)

               with torch.no_grad():
                  UpperCAmelCase__				:					Union[str, Any] 					=	pt_model_loaded(**lowerCamelCase_				).to_tuple()

               self.assertEqual(
                   len(lowerCamelCase_				)       ,len(lowerCamelCase_				)       ,'''Output lengths differ between Flax and PyTorch'''				)
               for fx_output, pt_output in zip(lowerCamelCase_       ,lowerCamelCase_				):
                  self.assert_almost_equals(fx_output[:, -1]       ,pt_output[:, -1].numpy()       ,4e-2				)




      @tooslow
      def 		lowerCAmelCase__       (					self				)      ->      List[Any]:



         '''simple docstring'''



         for model_class_name in self.all_model_classes:
            UpperCAmelCase__				:					Optional[Any] 					=	model_class_name.from_pretrained('''EleutherAI/gpt-j-6B'''				)
            UpperCAmelCase__				:					Optional[Any] 					=	model(np.ones((1, 1)				)				)
            self.assertIsNotNone(lowerCamelCase_				)





'''simple docstring'''

from __future__ import absolute_import, division, print_function, unicode_literals

from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers import RobertaConfig
from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.roberta.modeling_roberta import (
    ROBERTA_INPUTS_DOCSTRING,
    ROBERTA_START_DOCSTRING,
    RobertaEmbeddings,
)

from .modeling_highway_bert import BertPreTrainedModel, DeeBertModel, HighwayException, entropy


@add_start_docstrings(
    '''The RoBERTa Model transformer with early exiting (DeeRoBERTa). '''		,lowerCAmelCase		,)
class 					_lowercase						(   lowerCAmelCase ):
      '''simple docstring'''


      UpperCAmelCase_   :   int   				= RobertaConfig
      UpperCAmelCase_   :   Tuple   				= '''roberta'''




      def __init__(					self       ,lowerCamelCase_				)      ->      List[Any]:



         '''simple docstring'''



         super().__init__(lowerCamelCase_				)

         UpperCAmelCase__				:					Optional[Any] 					=	RobertaEmbeddings(lowerCamelCase_				)
         self.init_weights()




@add_start_docstrings(
    '''RoBERTa Model (with early exiting - DeeRoBERTa) with a classifier on top,
    also takes care of multi-layer training. '''		,lowerCAmelCase		,)
class 					_lowercase						(   lowerCAmelCase ):
      '''simple docstring'''


      UpperCAmelCase_   :   Tuple   				= RobertaConfig
      UpperCAmelCase_   :   Optional[Any]   				= '''roberta'''




      def __init__(					self       ,lowerCamelCase_				)      ->      str:



         '''simple docstring'''



         super().__init__(lowerCamelCase_				)
         UpperCAmelCase__				:					Any 					=	config.num_labels
         UpperCAmelCase__				:					Union[str, Any] 					=	config.num_hidden_layers

         UpperCAmelCase__				:					str 					=	DeeRobertaModel(lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[int] 					=	nn.Dropout(config.hidden_dropout_prob				)
         UpperCAmelCase__				:					Union[str, Any] 					=	nn.Linear(config.hidden_size       ,self.config.num_labels				)




      @add_start_docstrings_to_model_forward(lowerCamelCase_				)
      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,lowerCamelCase_=-1       ,lowerCamelCase_=False       ,)      ->      Optional[Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					Tuple 					=	self.num_layers
         try:
            UpperCAmelCase__				:					Union[str, Any] 					=	self.roberta(
                lowerCamelCase_       ,attention_mask=lowerCamelCase_       ,token_type_ids=lowerCamelCase_       ,position_ids=lowerCamelCase_       ,head_mask=lowerCamelCase_       ,inputs_embeds=lowerCamelCase_       ,)

            UpperCAmelCase__				:					List[Any] 					=	outputs[1]

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

         if not self.training:
            UpperCAmelCase__				:					Dict 					=	entropy(lowerCamelCase_				)
            UpperCAmelCase__				:					List[Any] 					=	[]
            UpperCAmelCase__				:					str 					=	[]
         if labels is not None:
            if self.num_labels == 1:
               #  We are doing regression
               UpperCAmelCase__				:					str 					=	MSELoss()
               UpperCAmelCase__				:					Union[str, Any] 					=	loss_fct(logits.view(-1				)       ,labels.view(-1				)				)
            else:
               UpperCAmelCase__				:					Any 					=	CrossEntropyLoss()
               UpperCAmelCase__				:					Optional[Any] 					=	loss_fct(logits.view(-1       ,self.num_labels				)       ,labels.view(-1				)				)

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

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

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




