






'''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 				lowercase       :


  '''simple docstring'''

  def __init__(					self     :     Union[str, Any]				,     __lowerCamelCase     :     List[str]				,     __lowerCamelCase     :     Tuple=13				,     __lowerCamelCase     :     Tuple=7				,     __lowerCamelCase     :     Any=True				,     __lowerCamelCase     :     List[Any]=True				,     __lowerCamelCase     :     Union[str, Any]=True				,     __lowerCamelCase     :     Any=True				,     __lowerCamelCase     :     List[str]=99				,     __lowerCamelCase     :     Any=64				,     __lowerCamelCase     :     Tuple=5				,     __lowerCamelCase     :     Tuple=4				,     __lowerCamelCase     :     Optional[int]=37				,     __lowerCamelCase     :     Union[str, Any]="gelu"				,     __lowerCamelCase     :     str=0.1				,     __lowerCamelCase     :     int=0.1				,     __lowerCamelCase     :     Union[str, Any]=512				,     __lowerCamelCase     :     Optional[int]=16				,     __lowerCamelCase     :     Optional[int]=2				,     __lowerCamelCase     :     Tuple=0.0_2				,     __lowerCamelCase     :     int=3				,     __lowerCamelCase     :     int=4				,     __lowerCamelCase     :     Optional[Any]=None				,     )							->							str:

   '''simple docstring'''




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

  def    a__				(					self     :     Any     )							->							Optional[Any]:

   '''simple docstring'''




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

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

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

   lowerCamelCase__  =     self.get_config()

   return config, input_ids, input_mask, token_labels

  def    a__				(					self     :     Any     )							->							Dict:

   '''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=__lowerCamelCase				,     initializer_range=self.initializer_range				,     pad_token_id=self.pad_token_id				,     )

  def    a__				(					self     :     int     )							->							Tuple:

   '''simple docstring'''




   lowerCamelCase__  ,	lowerCamelCase__  ,	lowerCamelCase__  ,	lowerCamelCase__  =     self.prepare_config_and_inputs()

   lowerCamelCase__  =     True

   return config, input_ids, input_mask, token_labels

  def    a__				(					self     :     int				,     __lowerCamelCase     :     List[str]				,     __lowerCamelCase     :     Union[str, Any]				,     __lowerCamelCase     :     str     )							->							List[str]:

   '''simple docstring'''




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

  def    a__				(					self     :     Union[str, Any]				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     int     )							->							int:

   '''simple docstring'''




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

  def    a__				(					self     :     int				,     __lowerCamelCase     :     Union[str, Any]				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     int     )							->							List[Any]:

   '''simple docstring'''




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

  def    a__				(					self     :     Optional[int]				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     Union[str, Any]				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     List[Any]     )							->							int:

   '''simple docstring'''




   lowerCamelCase__  =     self.num_labels
   lowerCamelCase__  =     GPTNeoXForQuestionAnswering(__lowerCamelCase     )
   model.to(__lowerCamelCase     )
   model.eval()
   lowerCamelCase__  =     model(__lowerCamelCase				,     attention_mask=__lowerCamelCase     )
   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    a__				(					self     :     List[str]				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     List[str]				,     __lowerCamelCase     :     List[Any]     )							->							Optional[int]:

   '''simple docstring'''




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

  def    a__				(					self     :     str				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     Optional[int]     )							->							Tuple:

   '''simple docstring'''




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

  def    a__				(					self     :     Tuple				,     __lowerCamelCase     :     Union[str, Any]				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     List[str]     )							->							Tuple:

   '''simple docstring'''




   lowerCamelCase__  =     True
   lowerCamelCase__  =     GPTNeoXForCausalLM(config=__lowerCamelCase     )
   model.to(__lowerCamelCase     )
   model.eval()

   # first forward pass
   lowerCamelCase__  =     model(__lowerCamelCase				,     attention_mask=__lowerCamelCase				,     use_cache=__lowerCamelCase     )
   lowerCamelCase__  =     outputs.past_key_values

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

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

   lowerCamelCase__  =     model(__lowerCamelCase				,     attention_mask=__lowerCamelCase				,     output_hidden_states=__lowerCamelCase     )
   lowerCamelCase__  =     output_from_no_past["hidden_states"][0]
   lowerCamelCase__  =     model(
       __lowerCamelCase				,     attention_mask=__lowerCamelCase				,     past_key_values=__lowerCamelCase				,     output_hidden_states=__lowerCamelCase				,     )["hidden_states"][0]

   # select random slice
   lowerCamelCase__  =     ids_tensor((1,)				,     output_from_past.shape[-1]     ).item()
   lowerCamelCase__  =     output_from_no_past[:, -3:, random_slice_idx].detach()
   lowerCamelCase__  =     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(__lowerCamelCase				,     __lowerCamelCase				,     atol=1E-3     )     )




  def    a__				(					self     :     Optional[int]     )							->							Any:

   '''simple docstring'''




   lowerCamelCase__  =     self.prepare_config_and_inputs()
   lowerCamelCase__  ,	lowerCamelCase__  ,	lowerCamelCase__  ,	lowerCamelCase__  =     config_and_inputs
   lowerCamelCase__  =     {"input_ids": input_ids, "attention_mask": input_mask}
   return config, inputs_dict
@require_torch
class 				lowercase       (							_lowerCamelCase ,       _lowerCamelCase ,       _lowerCamelCase ,       unittest.TestCase      ):


  '''simple docstring'''

  lowerCAmelCase__      				=							(
      (
          GPTNeoXModel,
          GPTNeoXForCausalLM,
          GPTNeoXForQuestionAnswering,
          GPTNeoXForSequenceClassification,
          GPTNeoXForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  lowerCAmelCase__      				=							(GPTNeoXForCausalLM,) if is_torch_available() else ()
  lowerCAmelCase__      				=							(
      {
          "feature-extraction": GPTNeoXModel,
          "question-answering": GPTNeoXForQuestionAnswering,
          "text-classification": GPTNeoXForSequenceClassification,
          "text-generation": GPTNeoXForCausalLM,
          "token-classification": GPTNeoXForTokenClassification,
          "zero-shot": GPTNeoXForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  lowerCAmelCase__      				=							False
  lowerCAmelCase__      				=							False
  lowerCAmelCase__      				=							False
  lowerCAmelCase__      				=							False

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

   '''simple docstring'''




   lowerCamelCase__  =     GPTNeoXModelTester(self     )
   lowerCamelCase__  =     ConfigTester(self				,     config_class=__lowerCamelCase				,     hidden_size=64				,     num_attention_heads=8     )

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

   '''simple docstring'''




   self.config_tester.run_common_tests()

  def    a__				(					self     :     Dict     )							->							Optional[Any]:

   '''simple docstring'''




   lowerCamelCase__  ,	lowerCamelCase__  ,	lowerCamelCase__  ,	lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(__lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase     )

  def    a__				(					self     :     Dict     )							->							List[str]:

   '''simple docstring'''




   lowerCamelCase__  ,	lowerCamelCase__  ,	lowerCamelCase__  ,	lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs_for_decoder()
   self.model_tester.create_and_check_model_as_decoder(__lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase     )

  def    a__				(					self     :     Any     )							->							int:

   '''simple docstring'''




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

   lowerCamelCase__  =     None

   self.model_tester.create_and_check_model_as_decoder(__lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase     )

  def    a__				(					self     :     List[str]     )							->							Tuple:

   '''simple docstring'''




   lowerCamelCase__  ,	lowerCamelCase__  ,	lowerCamelCase__  ,	lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_decoder_model_past_large_inputs(__lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase     )

  def    a__				(					self     :     List[Any]     )							->							Optional[Any]:

   '''simple docstring'''




   lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_causal_lm(*__lowerCamelCase     )

  def    a__				(					self     :     List[Any]     )							->							str:

   '''simple docstring'''




   lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_question_answering(*__lowerCamelCase     )

  def    a__				(					self     :     List[str]     )							->							Any:

   '''simple docstring'''




   lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_sequence_classification(*__lowerCamelCase     )

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

   '''simple docstring'''




   lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_token_classification(*__lowerCamelCase     )

  @unittest.skip(reason="Feed forward chunking is not implemented"     )
  def    a__				(					self     :     int     )							->							Dict:

   '''simple docstring'''




   pass




  @parameterized.expand([("linear",), ("dynamic",)]     )
  def    a__				(					self     :     str				,     __lowerCamelCase     :     Optional[int]     )							->							Dict:

   '''simple docstring'''




   lowerCamelCase__  ,	lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs_for_common()
   lowerCamelCase__  =     ids_tensor([1, 10]				,     config.vocab_size     )
   lowerCamelCase__  =     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
   lowerCamelCase__  =     GPTNeoXModel(__lowerCamelCase     )
   original_model.to(__lowerCamelCase     )
   original_model.eval()
   lowerCamelCase__  =     original_model(__lowerCamelCase     ).last_hidden_state
   lowerCamelCase__  =     original_model(__lowerCamelCase     ).last_hidden_state

   set_seed(42     )  # Fixed seed at init time so the two models get the same random weights
   lowerCamelCase__  =     {"type": scaling_type, "factor": 10.0}
   lowerCamelCase__  =     GPTNeoXModel(__lowerCamelCase     )
   scaled_model.to(__lowerCamelCase     )
   scaled_model.eval()
   lowerCamelCase__  =     scaled_model(__lowerCamelCase     ).last_hidden_state
   lowerCamelCase__  =     scaled_model(__lowerCamelCase     ).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(__lowerCamelCase				,     __lowerCamelCase				,     atol=1E-5     )     )
   else:
    self.assertFalse(torch.allclose(__lowerCamelCase				,     __lowerCamelCase				,     atol=1E-5     )     )

   # The output should be different for long inputs
   self.assertFalse(torch.allclose(__lowerCamelCase				,     __lowerCamelCase				,     atol=1E-5     )     )
@require_torch
class 				lowercase       (							unittest.TestCase      ):


  '''simple docstring'''

  @slow
  def    a__				(					self     :     Tuple     )							->							Dict:

   '''simple docstring'''




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

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

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

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

    self.assertEqual(__lowerCamelCase				,     __lowerCamelCase     )









'''simple docstring'''





import tensorflow as tf

from ...tf_utils import shape_list
class 				lowercase       (							tf.keras.layers.Layer      ):


    '''simple docstring'''

    def __init__(					self     :     Union[str, Any]				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     Union[str, Any]				,     __lowerCamelCase     :     int=1				,     __lowerCamelCase     :     Tuple=False				,     **__lowerCamelCase     :     Dict     )							->							str:

     '''simple docstring'''




     super().__init__(**__lowerCamelCase     )

     lowerCamelCase__  =     vocab_size
     lowerCamelCase__  =     d_embed
     lowerCamelCase__  =     d_proj

     lowerCamelCase__  =     cutoffs + [vocab_size]
     lowerCamelCase__  =     [0] + self.cutoffs
     lowerCamelCase__  =     div_val

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

     lowerCamelCase__  =     []
     lowerCamelCase__  =     []

    def    a__				(					self     :     Optional[int]				,     __lowerCamelCase     :     str     )							->							List[str]:

     '''simple docstring'''




     if self.n_clusters > 0:
      lowerCamelCase__  =     self.add_weight(
          shape=(self.n_clusters, self.d_embed)				,     initializer="zeros"				,     trainable=__lowerCamelCase				,     name="cluster_weight"     )
      lowerCamelCase__  =     self.add_weight(
          shape=(self.n_clusters,)				,     initializer="zeros"				,     trainable=__lowerCamelCase				,     name="cluster_bias"     )

     if self.div_val == 1:
      for i in range(len(self.cutoffs     )     ):
       if self.d_proj != self.d_embed:
        lowerCamelCase__  =     self.add_weight(
            shape=(self.d_embed, self.d_proj)				,     initializer="zeros"				,     trainable=__lowerCamelCase				,     name=f'''out_projs_._{i}'''				,     )
        self.out_projs.append(__lowerCamelCase     )
       else:
        self.out_projs.append(__lowerCamelCase     )
       lowerCamelCase__  =     self.add_weight(
           shape=(self.vocab_size, self.d_embed)				,     initializer="zeros"				,     trainable=__lowerCamelCase				,     name=f'''out_layers_._{i}_._weight'''				,     )
       lowerCamelCase__  =     self.add_weight(
           shape=(self.vocab_size,)				,     initializer="zeros"				,     trainable=__lowerCamelCase				,     name=f'''out_layers_._{i}_._bias'''				,     )
       self.out_layers.append((weight, bias)     )
     else:
      for i in range(len(self.cutoffs     )     ):
       lowerCamelCase__  ,	lowerCamelCase__  =     self.cutoff_ends[i], self.cutoff_ends[i + 1]
       lowerCamelCase__  =     self.d_embed // (self.div_val**i)

       lowerCamelCase__  =     self.add_weight(
           shape=(d_emb_i, self.d_proj)				,     initializer="zeros"				,     trainable=__lowerCamelCase				,     name=f'''out_projs_._{i}'''     )
       self.out_projs.append(__lowerCamelCase     )
       lowerCamelCase__  =     self.add_weight(
           shape=(r_idx - l_idx, d_emb_i)				,     initializer="zeros"				,     trainable=__lowerCamelCase				,     name=f'''out_layers_._{i}_._weight'''				,     )
       lowerCamelCase__  =     self.add_weight(
           shape=(r_idx - l_idx,)				,     initializer="zeros"				,     trainable=__lowerCamelCase				,     name=f'''out_layers_._{i}_._bias'''				,     )
       self.out_layers.append((weight, bias)     )
     super().build(__lowerCamelCase     )

    @staticmethod
    def    a__				(					__lowerCamelCase     :     Union[str, Any]				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     List[str]=None     )							->							str:

     '''simple docstring'''




     lowerCamelCase__  =     x
     if proj is not None:
      lowerCamelCase__  =     tf.einsum("ibd,ed->ibe"				,     __lowerCamelCase				,     __lowerCamelCase     )
     return tf.einsum("ibd,nd->ibn"				,     __lowerCamelCase				,     __lowerCamelCase     ) + b

    @staticmethod
    def    a__				(					__lowerCamelCase     :     List[str]				,     __lowerCamelCase     :     Union[str, Any]     )							->							Union[str, Any]:

     '''simple docstring'''




     lowerCamelCase__  =     shape_list(__lowerCamelCase     )
     lowerCamelCase__  =     tf.range(lp_size[0]				,     dtype=target.dtype     )
     lowerCamelCase__  =     tf.stack([r, target]				,     1     )
     return tf.gather_nd(__lowerCamelCase				,     __lowerCamelCase     )




    def    a__				(					self     :     Any				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     str=True				,     __lowerCamelCase     :     Tuple=False     )							->							int:

     '''simple docstring'''




     lowerCamelCase__  =     0
     if self.n_clusters == 0:
      lowerCamelCase__  =     self._logit(__lowerCamelCase				,     self.out_layers[0][0]				,     self.out_layers[0][1]				,     self.out_projs[0]     )
      if target is not None:
       lowerCamelCase__  =     tf.nn.sparse_softmax_cross_entropy_with_logits(labels=__lowerCamelCase				,     logits=__lowerCamelCase     )
      lowerCamelCase__  =     tf.nn.log_softmax(__lowerCamelCase				,     axis=-1     )
     else:
      lowerCamelCase__  =     shape_list(__lowerCamelCase     )
      lowerCamelCase__  =     []
      lowerCamelCase__  =     tf.zeros(hidden_sizes[:2]     )
      for i in range(len(self.cutoffs     )     ):
       lowerCamelCase__  ,	lowerCamelCase__  =     self.cutoff_ends[i], self.cutoff_ends[i + 1]
       if target is not None:
        lowerCamelCase__  =     (target >= l_idx) & (target < r_idx)
        lowerCamelCase__  =     tf.where(__lowerCamelCase     )
        lowerCamelCase__  =     tf.boolean_mask(__lowerCamelCase				,     __lowerCamelCase     ) - l_idx

       if self.div_val == 1:
        lowerCamelCase__  =     self.out_layers[0][0][l_idx:r_idx]
        lowerCamelCase__  =     self.out_layers[0][1][l_idx:r_idx]
       else:
        lowerCamelCase__  =     self.out_layers[i][0]
        lowerCamelCase__  =     self.out_layers[i][1]

       if i == 0:
        lowerCamelCase__  =     tf.concat([cur_W, self.cluster_weight]				,     0     )
        lowerCamelCase__  =     tf.concat([cur_b, self.cluster_bias]				,     0     )

        lowerCamelCase__  =     self._logit(__lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase				,     self.out_projs[0]     )
        lowerCamelCase__  =     tf.nn.log_softmax(__lowerCamelCase     )
        out.append(head_logprob[..., : self.cutoffs[0]]     )
        if target is not None:
         lowerCamelCase__  =     tf.boolean_mask(__lowerCamelCase				,     __lowerCamelCase     )
         lowerCamelCase__  =     self._gather_logprob(__lowerCamelCase				,     __lowerCamelCase     )
       else:
        lowerCamelCase__  =     self._logit(__lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase				,     self.out_projs[i]     )
        lowerCamelCase__  =     tf.nn.log_softmax(__lowerCamelCase     )
        lowerCamelCase__  =     self.cutoffs[0] + i - 1  # No probability for the head cluster
        lowerCamelCase__  =     head_logprob[..., cluster_prob_idx, None] + tail_logprob
        out.append(__lowerCamelCase     )
        if target is not None:
         lowerCamelCase__  =     tf.boolean_mask(__lowerCamelCase				,     __lowerCamelCase     )
         lowerCamelCase__  =     tf.boolean_mask(__lowerCamelCase				,     __lowerCamelCase     )
         lowerCamelCase__  =     self._gather_logprob(__lowerCamelCase				,     __lowerCamelCase     )
         cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
       if target is not None:
        loss += tf.scatter_nd(__lowerCamelCase				,     -cur_logprob				,     shape_list(__lowerCamelCase     )     )
      lowerCamelCase__  =     tf.concat(__lowerCamelCase				,     axis=-1     )

     if target is not None:
      if return_mean:
       lowerCamelCase__  =     tf.reduce_mean(__lowerCamelCase     )
      # Add the training-time loss value to the layer using `self.add_loss()`.
      self.add_loss(__lowerCamelCase     )

      # Log the loss as a metric (we could log arbitrary metrics,
      # including different metrics for training and inference.
      self.add_metric(__lowerCamelCase				,     name=self.name				,     aggregation="mean" if return_mean else ""     )

     return out

