"""simple docstring"""
import unittest

from parameterized import parameterized

from transformers import AutoTokenizer, GPTNeoXConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import (
					    GPTNeoXForCausalLM,
					    GPTNeoXForQuestionAnswering,
					    GPTNeoXForSequenceClassification,
					    GPTNeoXForTokenClassification,
					    GPTNeoXModel,
					)



class      lowerCAmelCase_			:



				"""simple docstring"""

				def __init__( self						,			lowerCAmelCase						,			lowerCAmelCase=13						,			lowerCAmelCase=7						,			lowerCAmelCase=True						,			lowerCAmelCase=True						,			lowerCAmelCase=True						,			lowerCAmelCase=True						,			lowerCAmelCase=99						,			lowerCAmelCase=64						,			lowerCAmelCase=5						,			lowerCAmelCase=4						,			lowerCAmelCase=37						,			lowerCAmelCase="gelu"						,			lowerCAmelCase=0.1						,			lowerCAmelCase=0.1						,			lowerCAmelCase=5_12						,			lowerCAmelCase=16						,			lowerCAmelCase=2						,			lowerCAmelCase=0.02						,			lowerCAmelCase=3						,			lowerCAmelCase=4						,			lowerCAmelCase=None						,			):






										"""simple docstring"""

										snake_case	    =	parent
										snake_case	    =	batch_size
										snake_case	    =	seq_length
										snake_case	    =	is_training
										snake_case	    =	use_input_mask
										snake_case	    =	use_token_type_ids
										snake_case	    =	use_labels
										snake_case	    =	vocab_size
										snake_case	    =	hidden_size
										snake_case	    =	num_hidden_layers
										snake_case	    =	num_attention_heads
										snake_case	    =	intermediate_size
										snake_case	    =	hidden_act
										snake_case	    =	hidden_dropout_prob
										snake_case	    =	attention_probs_dropout_prob
										snake_case	    =	max_position_embeddings
										snake_case	    =	type_vocab_size
										snake_case	    =	type_sequence_label_size
										snake_case	    =	initializer_range
										snake_case	    =	num_labels
										snake_case	    =	num_choices
										snake_case	    =	scope
										snake_case	    =	vocab_size - 1

				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	ids_tensor([self.batch_size, self.seq_length]						,			self.vocab_size			)

										snake_case	    =	None
										if self.use_input_mask:
																snake_case	    =	random_attention_mask([self.batch_size, self.seq_length]			)

										snake_case	    =	None
										if self.use_labels:
																snake_case	    =	ids_tensor([self.batch_size, self.seq_length]						,			self.num_labels			)

										snake_case	    =	self.get_config()

										return config, input_ids, input_mask, token_labels

				def 						snake_case       ( self			):






										"""simple docstring"""

										return GPTNeoXConfig(
										    vocab_size=self.vocab_size						,			hidden_size=self.hidden_size						,			num_hidden_layers=self.num_hidden_layers						,			num_attention_heads=self.num_attention_heads						,			intermediate_size=self.intermediate_size						,			hidden_act=self.hidden_act						,			hidden_dropout_prob=self.hidden_dropout_prob						,			attention_probs_dropout_prob=self.attention_probs_dropout_prob						,			max_position_embeddings=self.max_position_embeddings						,			type_vocab_size=self.type_vocab_size						,			is_decoder=lowercase__						,			initializer_range=self.initializer_range						,			pad_token_id=self.pad_token_id						,			)

				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case    ,snake_case    ,snake_case    ,snake_case	    =	self.prepare_config_and_inputs()

										snake_case	    =	True

										return config, input_ids, input_mask, token_labels

				def 						snake_case       ( self						,			lowerCAmelCase						,			lowerCAmelCase						,			lowerCAmelCase			):






										"""simple docstring"""

										snake_case	    =	GPTNeoXModel(config=lowercase__			)
										model.to(lowercase__			)
										model.eval()
										snake_case	    =	model(lowercase__						,			attention_mask=lowercase__			)
										snake_case	    =	model(lowercase__			)
										self.parent.assertEqual(result.last_hidden_state.shape						,			(self.batch_size, self.seq_length, self.hidden_size)			)

				def 						snake_case       ( self						,			lowerCAmelCase						,			lowerCAmelCase						,			lowerCAmelCase			):






										"""simple docstring"""

										snake_case	    =	True
										snake_case	    =	GPTNeoXModel(lowercase__			)
										model.to(lowercase__			)
										model.eval()
										snake_case	    =	model(lowercase__						,			attention_mask=lowercase__			)
										self.parent.assertEqual(result.last_hidden_state.shape						,			(self.batch_size, self.seq_length, self.hidden_size)			)

				def 						snake_case       ( self						,			lowerCAmelCase						,			lowerCAmelCase						,			lowerCAmelCase						,			lowerCAmelCase			):






										"""simple docstring"""

										snake_case	    =	GPTNeoXForCausalLM(config=lowercase__			)
										model.to(lowercase__			)
										model.eval()
										snake_case	    =	model(lowercase__						,			attention_mask=lowercase__						,			labels=lowercase__			)
										self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.seq_length, self.vocab_size)			)

				def 						snake_case       ( self						,			lowerCAmelCase						,			lowerCAmelCase						,			lowerCAmelCase						,			lowerCAmelCase			):






										"""simple docstring"""

										snake_case	    =	self.num_labels
										snake_case	    =	GPTNeoXForQuestionAnswering(lowercase__			)
										model.to(lowercase__			)
										model.eval()
										snake_case	    =	model(lowercase__						,			attention_mask=lowercase__			)
										self.parent.assertEqual(result.start_logits.shape						,			(self.batch_size, self.seq_length)			)
										self.parent.assertEqual(result.end_logits.shape						,			(self.batch_size, self.seq_length)			)

				def 						snake_case       ( self						,			lowerCAmelCase						,			lowerCAmelCase						,			lowerCAmelCase						,			lowerCAmelCase			):






										"""simple docstring"""

										snake_case	    =	self.num_labels
										snake_case	    =	GPTNeoXForSequenceClassification(lowercase__			)
										model.to(lowercase__			)
										model.eval()
										snake_case	    =	ids_tensor([self.batch_size]						,			self.type_sequence_label_size			)
										snake_case	    =	model(lowercase__						,			attention_mask=lowercase__						,			labels=lowercase__			)
										self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.num_labels)			)

				def 						snake_case       ( self						,			lowerCAmelCase						,			lowerCAmelCase						,			lowerCAmelCase						,			lowerCAmelCase			):






										"""simple docstring"""

										snake_case	    =	self.num_labels
										snake_case	    =	GPTNeoXForTokenClassification(lowercase__			)
										model.to(lowercase__			)
										model.eval()
										snake_case	    =	model(lowercase__						,			attention_mask=lowercase__						,			labels=lowercase__			)
										self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.seq_length, self.num_labels)			)

				def 						snake_case       ( self						,			lowerCAmelCase						,			lowerCAmelCase						,			lowerCAmelCase			):






										"""simple docstring"""

										snake_case	    =	True
										snake_case	    =	GPTNeoXForCausalLM(config=lowercase__			)
										model.to(lowercase__			)
										model.eval()

										# first forward pass
										snake_case	    =	model(lowercase__						,			attention_mask=lowercase__						,			use_cache=lowercase__			)
										snake_case	    =	outputs.past_key_values

										# create hypothetical multiple next token and extent to next_input_ids
										snake_case	    =	ids_tensor((self.batch_size, 3)						,			config.vocab_size			)
										snake_case	    =	ids_tensor((self.batch_size, 3)						,			vocab_size=2			)

										# append to next input_ids and
										snake_case	    =	torch.cat([input_ids, next_tokens]						,			dim=-1			)
										snake_case	    =	torch.cat([input_mask, next_mask]						,			dim=-1			)

										snake_case	    =	model(lowercase__						,			attention_mask=lowercase__						,			output_hidden_states=lowercase__			)
										snake_case	    =	output_from_no_past['hidden_states'][0]
										snake_case	    =	model(
										    lowercase__						,			attention_mask=lowercase__						,			past_key_values=lowercase__						,			output_hidden_states=lowercase__						,			)['hidden_states'][0]

										# select random slice
										snake_case	    =	ids_tensor((1,)						,			output_from_past.shape[-1]			).item()
										snake_case	    =	output_from_no_past[:, -3:, random_slice_idx].detach()
										snake_case	    =	output_from_past[:, :, random_slice_idx].detach()

										self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]			)

										# test that outputs are equal for slice
										self.parent.assertTrue(torch.allclose(lowercase__						,			lowercase__						,			atol=1E-3			)			)

				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	self.prepare_config_and_inputs()
										snake_case    ,snake_case    ,snake_case    ,snake_case	    =	config_and_inputs
										snake_case	    =	{'input_ids': input_ids, 'attention_mask': input_mask}
										return config, inputs_dict





@require_torch
class      lowerCAmelCase_			(	lowerCamelCase__   ,			lowerCamelCase__   ,			lowerCamelCase__   ,			unittest.TestCase    ):



				"""simple docstring"""
				_lowerCAmelCase   :			Optional[Any]	     =  (
				    (
				        GPTNeoXModel,
				        GPTNeoXForCausalLM,
				        GPTNeoXForQuestionAnswering,
				        GPTNeoXForSequenceClassification,
				        GPTNeoXForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				_lowerCAmelCase   :			Union[str, Any]	     =  (GPTNeoXForCausalLM,) if is_torch_available() else ()
				_lowerCAmelCase   :			Union[str, Any]	     =  (
				    {
				        'feature-extraction': GPTNeoXModel,
				        'question-answering': GPTNeoXForQuestionAnswering,
				        'text-classification': GPTNeoXForSequenceClassification,
				        'text-generation': GPTNeoXForCausalLM,
				        'token-classification': GPTNeoXForTokenClassification,
				        'zero-shot': GPTNeoXForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				_lowerCAmelCase   :			List[str]	     =  False
				_lowerCAmelCase   :			Optional[int]	     =  False
				_lowerCAmelCase   :			List[str]	     =  False
				_lowerCAmelCase   :			Dict	     =  False

				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	GPTNeoXModelTester(self			)
										snake_case	    =	ConfigTester(self						,			config_class=lowercase__						,			hidden_size=64						,			num_attention_heads=8			)

				def 						snake_case       ( self			):






										"""simple docstring"""

										self.config_tester.run_common_tests()

				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case    ,snake_case    ,snake_case    ,snake_case	    =	self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(lowercase__						,			lowercase__						,			lowercase__			)

				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case    ,snake_case    ,snake_case    ,snake_case	    =	self.model_tester.prepare_config_and_inputs_for_decoder()
										self.model_tester.create_and_check_model_as_decoder(lowercase__						,			lowercase__						,			lowercase__			)

				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case    ,snake_case    ,snake_case    ,snake_case	    =	self.model_tester.prepare_config_and_inputs_for_decoder()

										snake_case	    =	None

										self.model_tester.create_and_check_model_as_decoder(lowercase__						,			lowercase__						,			lowercase__			)

				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case    ,snake_case    ,snake_case    ,snake_case	    =	self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_decoder_model_past_large_inputs(lowercase__						,			lowercase__						,			lowercase__			)

				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_causal_lm(*lowercase__			)

				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_question_answering(*lowercase__			)

				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_sequence_classification(*lowercase__			)

				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*lowercase__			)

				@unittest.skip(reason='Feed forward chunking is not implemented'			)
				def 						snake_case       ( self			):






										"""simple docstring"""

										pass

				@parameterized.expand([('linear',), ('dynamic',)]			)
				def 						snake_case       ( self						,			lowerCAmelCase			):






										"""simple docstring"""

										snake_case    ,snake_case	    =	self.model_tester.prepare_config_and_inputs_for_common()
										snake_case	    =	ids_tensor([1, 10]						,			config.vocab_size			)
										snake_case	    =	ids_tensor([1, int(config.max_position_embeddings * 1.5			)]						,			config.vocab_size			)

										set_seed(42			)  # Fixed seed at init time so the two models get the same random weights
										snake_case	    =	GPTNeoXModel(lowercase__			)
										original_model.to(lowercase__			)
										original_model.eval()
										snake_case	    =	original_model(lowercase__			).last_hidden_state
										snake_case	    =	original_model(lowercase__			).last_hidden_state

										set_seed(42			)  # Fixed seed at init time so the two models get the same random weights
										snake_case	    =	{'type': scaling_type, 'factor': 10.0}
										snake_case	    =	GPTNeoXModel(lowercase__			)
										scaled_model.to(lowercase__			)
										scaled_model.eval()
										snake_case	    =	scaled_model(lowercase__			).last_hidden_state
										snake_case	    =	scaled_model(lowercase__			).last_hidden_state

										# Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
										# maximum sequence length, so the outputs for the short input should match.
										if scaling_type == "dynamic":
																self.assertTrue(torch.allclose(lowercase__						,			lowercase__						,			atol=1E-5			)			)
										else:
																self.assertFalse(torch.allclose(lowercase__						,			lowercase__						,			atol=1E-5			)			)

										# The output should be different for long inputs
										self.assertFalse(torch.allclose(lowercase__						,			lowercase__						,			atol=1E-5			)			)





@require_torch
class      lowerCAmelCase_			(	unittest.TestCase    ):



				"""simple docstring"""

				@slow
				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	AutoTokenizer.from_pretrained('EleutherAI/pythia-410m-deduped'			)
										for checkpointing in [True, False]:
																snake_case	    =	GPTNeoXForCausalLM.from_pretrained('EleutherAI/pythia-410m-deduped'			)

																if checkpointing:
																						model.gradient_checkpointing_enable()
																else:
																						model.gradient_checkpointing_disable()
																model.to(lowercase__			)

																snake_case	    =	tokenizer('My favorite food is'						,			return_tensors='pt'			).to(lowercase__			)
																# The hub repo. is updated on 2023-04-04, resulting in poor outputs.
																# See: https://github.com/huggingface/transformers/pull/24193
																snake_case	    =	'My favorite food is a good old-fashioned, old-fashioned, old-fashioned.\n\nI\'m not sure'

																snake_case	    =	model.generate(**lowercase__						,			do_sample=lowercase__						,			max_new_tokens=20			)
																snake_case	    =	tokenizer.batch_decode(lowercase__			)[0]

																self.assertEqual(lowercase__						,			lowercase__			)





'''simple docstring'''






import torch
from torch import nn






class 			lowercase_	(nn.Module      ):




  """simple docstring"""




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

   __lowercase							  = n_token
   __lowercase							  = d_embed
   __lowercase							  = d_proj

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

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

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

   __lowercase							  = nn.ModuleList()
   __lowercase							  = nn.ParameterList()

   if div_val == 1:
    for i in range(len(self.cutoffs			)			):
     if d_proj != d_embed:
      self.out_projs.append(nn.Parameter(torch.FloatTensor(lowercase__					,lowercase__			)			)			)
     else:
      self.out_projs.append(lowercase__			)

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

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

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

   __lowercase							  = keep_order






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

   return logit






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

   if labels is not None:
    # Shift so that tokens < n predict n
    __lowercase							  = hidden[..., :-1, :].contiguous()
    __lowercase							  = labels[..., 1:].contiguous()
    __lowercase							  = hidden.view(-1					,hidden.size(-1			)			)
    __lowercase							  = labels.view(-1			)
    if hidden.size(0			) != labels.size(0			):
     raise RuntimeError('''Input and labels should have the same size in the batch dimension.'''			)
   else:
    __lowercase							  = hidden.view(-1					,hidden.size(-1			)			)

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

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

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

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

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

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

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

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

      if indices_i.numel() == 0:
       continue

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

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

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

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

   return out






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

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

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

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

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

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

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

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

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

    return out

