

'''simple docstring'''







from __future__ import annotations

import unittest

from transformers import RoFormerConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
      import tensorflow as tf

      from transformers import (
          TFRoFormerForCausalLM,
          TFRoFormerForMaskedLM,
          TFRoFormerForMultipleChoice,
          TFRoFormerForQuestionAnswering,
          TFRoFormerForSequenceClassification,
          TFRoFormerForTokenClassification,
          TFRoFormerModel,
      )
      from transformers.models.roformer.modeling_tf_roformer import (
          TFRoFormerSelfAttention,
          TFRoFormerSinusoidalPositionalEmbedding,
      )


class 					A__      :



    def __init__(					self		,				UpperCamelCase__		,				UpperCamelCase__=13		,				UpperCamelCase__=7		,				UpperCamelCase__=True		,				UpperCamelCase__=True		,				UpperCamelCase__=True		,				UpperCamelCase__=True		,				UpperCamelCase__=99		,				UpperCamelCase__=32		,				UpperCamelCase__=2		,				UpperCamelCase__=4		,				UpperCamelCase__=37		,				UpperCamelCase__="gelu"		,				UpperCamelCase__=0.1		,				UpperCamelCase__=0.1		,				UpperCamelCase__=512		,				UpperCamelCase__=16		,				UpperCamelCase__=2		,				UpperCamelCase__=0.02		,				UpperCamelCase__=3		,				UpperCamelCase__=4		,				UpperCamelCase__=None		,				)						->      int:





           '''simple docstring'''
           A_						=      parent
           A_						=      13
           A_						=      7
           A_						=      True
           A_						=      True
           A_						=      True
           A_						=      True
           A_						=      99
           A_						=      32
           A_						=      2
           A_						=      4
           A_						=      37
           A_						=      """gelu"""
           A_						=      0.1
           A_						=      0.1
           A_						=      512
           A_						=      16
           A_						=      2
           A_						=      0.02
           A_						=      3
           A_						=      4
           A_						=      None



    def 					snake_case_			(					self       )						->      Tuple:





           '''simple docstring'''
           A_						=      ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size       )

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

           A_						=      None
           if self.use_token_type_ids:
                  A_						=      ids_tensor([self.batch_size, self.seq_length]		,				self.type_vocab_size       )

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

           A_						=      RoFormerConfig(
               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		,				initializer_range=self.initializer_range		,				return_dict=A__		,				)

           return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



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





           '''simple docstring'''
           A_						=      TFRoFormerModel(config=A__       )
           A_						=      {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}

           A_						=      [input_ids, input_mask]
           A_						=      model(A__       )

           A_						=      model(A__       )

           self.parent.assertEqual(result.last_hidden_state.shape		,				(self.batch_size, self.seq_length, self.hidden_size)       )



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__       )						->      str:





           '''simple docstring'''
           A_						=      True
           A_						=      TFRoFormerForCausalLM(config=A__       )
           A_						=      {
               """input_ids""": input_ids,
               """attention_mask""": input_mask,
               """token_type_ids""": token_type_ids,
           }
           A_						=      model(A__       )["""logits"""]
           self.parent.assertListEqual(
               list(prediction_scores.numpy().shape       )		,				[self.batch_size, self.seq_length, self.vocab_size]       )



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__       )						->      Dict:





           '''simple docstring'''
           A_						=      TFRoFormerForMaskedLM(config=A__       )
           A_						=      {
               """input_ids""": input_ids,
               """attention_mask""": input_mask,
               """token_type_ids""": token_type_ids,
           }
           A_						=      model(A__       )
           self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.vocab_size)       )



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__       )						->      List[str]:





           '''simple docstring'''
           A_						=      self.num_labels
           A_						=      TFRoFormerForSequenceClassification(config=A__       )
           A_						=      {
               """input_ids""": input_ids,
               """attention_mask""": input_mask,
               """token_type_ids""": token_type_ids,
           }

           A_						=      model(A__       )
           self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.num_labels)       )



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__       )						->      Any:





           '''simple docstring'''
           A_						=      self.num_choices
           A_						=      TFRoFormerForMultipleChoice(config=A__       )
           A_						=      tf.tile(tf.expand_dims(A__		,				1       )		,				(1, self.num_choices, 1)       )
           A_						=      tf.tile(tf.expand_dims(A__		,				1       )		,				(1, self.num_choices, 1)       )
           A_						=      tf.tile(tf.expand_dims(A__		,				1       )		,				(1, self.num_choices, 1)       )
           A_						=      {
               """input_ids""": multiple_choice_inputs_ids,
               """attention_mask""": multiple_choice_input_mask,
               """token_type_ids""": multiple_choice_token_type_ids,
           }
           A_						=      model(A__       )
           self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.num_choices)       )



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__       )						->      List[str]:





           '''simple docstring'''
           A_						=      self.num_labels
           A_						=      TFRoFormerForTokenClassification(config=A__       )
           A_						=      {
               """input_ids""": input_ids,
               """attention_mask""": input_mask,
               """token_type_ids""": token_type_ids,
           }
           A_						=      model(A__       )
           self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.num_labels)       )



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__       )						->      Tuple:





           '''simple docstring'''
           A_						=      TFRoFormerForQuestionAnswering(config=A__       )
           A_						=      {
               """input_ids""": input_ids,
               """attention_mask""": input_mask,
               """token_type_ids""": token_type_ids,
           }

           A_						=      model(A__       )
           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       )						->      List[str]:





           '''simple docstring'''
           A_						=      self.prepare_config_and_inputs()
           (
               (
               A_							
           )					,			(
               A_							
           )					,			(
               A_							
           )					,			(
               A_							
           )					,			(
               A_							
           )					,			(
               A_							
           )					,			(
               A_							
           )					,										
           )						=      config_and_inputs
           A_						=      {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
           return config, inputs_dict


@require_tf
class 					A__      (      _lowerCamelCase						,						_lowerCamelCase						,						unittest.TestCase			):
    snake_case						=		(
        (
            TFRoFormerModel,
            TFRoFormerForCausalLM,
            TFRoFormerForMaskedLM,
            TFRoFormerForQuestionAnswering,
            TFRoFormerForSequenceClassification,
            TFRoFormerForTokenClassification,
            TFRoFormerForMultipleChoice,
        )
        if is_tf_available()
        else ()
    )
    snake_case						=		(
        {
            "feature-extraction": TFRoFormerModel,
            "fill-mask": TFRoFormerForMaskedLM,
            "question-answering": TFRoFormerForQuestionAnswering,
            "text-classification": TFRoFormerForSequenceClassification,
            "text-generation": TFRoFormerForCausalLM,
            "token-classification": TFRoFormerForTokenClassification,
            "zero-shot": TFRoFormerForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )

    snake_case						=		False
    snake_case						=		False



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__       )						->      str:





           '''simple docstring'''
           if pipeline_test_casse_name == "TextGenerationPipelineTests":
                  return True

           return False



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





           '''simple docstring'''
           A_						=      TFRoFormerModelTester(self       )
           A_						=      ConfigTester(self		,				config_class=A__		,				hidden_size=37       )



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





           '''simple docstring'''
           self.config_tester.run_common_tests()



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





           '''simple docstring'''
           A_						=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*A__       )



    def 					snake_case_			(					self       )						->      str:





           '''simple docstring'''
           A_						=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*A__       )



    def 					snake_case_			(					self       )						->      int:





           '''simple docstring'''
           A_						=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_lm_head(*A__       )



    def 					snake_case_			(					self       )						->      Dict:





           '''simple docstring'''
           A_						=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_multiple_choice(*A__       )



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





           '''simple docstring'''
           A_						=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(*A__       )



    def 					snake_case_			(					self       )						->      Tuple:





           '''simple docstring'''
           A_						=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(*A__       )



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





           '''simple docstring'''
           A_						=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*A__       )







    @slow
    def 					snake_case_			(					self       )						->      str:





           '''simple docstring'''
           A_						=      TFRoFormerModel.from_pretrained("""junnyu/roformer_chinese_base"""       )
           self.assertIsNotNone(A__       )


@require_tf
class 					A__      (      unittest.TestCase			):



    @slow
    def 					snake_case_			(					self       )						->      Optional[Any]:





           '''simple docstring'''
           A_						=      TFRoFormerForMaskedLM.from_pretrained("""junnyu/roformer_chinese_base"""       )
           A_						=      tf.constant([[0, 1, 2, 3, 4, 5]]       )
           A_						=      model(A__       )[0]

           # TODO Replace vocab size
           A_						=      50000

           A_						=      [1, 6, vocab_size]
           self.assertEqual(output.shape		,				A__       )

           print(output[:, :3, :3]       )

           # TODO Replace values below with what was printed above.
           A_						=      tf.constant(
               [
                   [
                       [-0.12053341, -1.0264901, 0.29221946],
                       [-1.5133783, 0.197433, 0.15190607],
                       [-5.0135403, -3.900256, -0.84038764],
                   ]
               ]       )
           tf.debugging.assert_near(output[:, :3, :3]		,				A__		,				atol=1e-4       )


@require_tf
class 					A__      (      unittest.TestCase			):
    snake_case						=		1e-4



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





           '''simple docstring'''
           A_						=      tf.constant([[4, 10]]       )
           A_						=      TFRoFormerSinusoidalPositionalEmbedding(num_positions=6		,				embedding_dim=6       )

           A_						=      emba(input_ids.shape       )
           A_						=      tf.constant(
               [[0.0000, 0.0000, 0.0000, 1.0000, 1.0000, 1.0000], [0.8415, 0.0464, 0.0022, 0.5403, 0.9989, 1.0000]]       )

           tf.debugging.assert_near(A__		,				A__		,				atol=self.tolerance       )







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





           '''simple docstring'''
           A_						=      tf.constant(
               [
                   [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
                   [0.8415, 0.8219, 0.8020, 0.7819, 0.7617],
                   [0.9093, 0.9364, 0.9581, 0.9749, 0.9870],
               ]       )
           A_						=      TFRoFormerSinusoidalPositionalEmbedding(num_positions=512		,				embedding_dim=512       )
           emba([2, 16, 512]       )
           A_						=      emba.weight[:3, :5]

           tf.debugging.assert_near(A__		,				A__		,				atol=self.tolerance       )


@require_tf
class 					A__      (      unittest.TestCase			):
    snake_case						=		1e-4



    def 					snake_case_			(					self       )						->      int:





           '''simple docstring'''
           # 2,12,16,64
           A_						=      tf.reshape(tf.range(2 * 12 * 16 * 64		,				dtype=tf.floataa       )		,				shape=(2, 12, 16, 64)       ) / 100

           A_						=      -tf.reshape(tf.range(2 * 12 * 16 * 64		,				dtype=tf.floataa       )		,				shape=(2, 12, 16, 64)       ) / 100

           A_						=      TFRoFormerSinusoidalPositionalEmbedding(num_positions=32		,				embedding_dim=64       )
           A_						=      embed_positions([2, 16, 768]       )[None, None, :, :]

           A_					,			A_						=      TFRoFormerSelfAttention.apply_rotary_position_embeddings(
               A__		,				A__		,				A__       )

           A_						=      tf.constant(
               [
                   [0.0000, 0.0100, 0.0200, 0.0300, 0.0400, 0.0500, 0.0600, 0.0700],
                   [-0.2012, 0.8897, 0.0263, 0.9401, 0.2074, 0.9463, 0.3481, 0.9343],
                   [-1.7057, 0.6271, -1.2145, 1.3897, -0.6303, 1.7647, -0.1173, 1.8985],
                   [-2.1731, -1.6397, -2.7358, 0.2854, -2.1840, 1.7183, -1.3018, 2.4871],
                   [0.2717, -3.6173, -2.9206, -2.1988, -3.6638, 0.3858, -2.9155, 2.2980],
                   [3.9859, -2.1580, -0.7984, -4.4904, -4.1181, -2.0252, -4.4782, 1.1253],
               ]       )
           A_						=      tf.constant(
               [
                   [0.0000, -0.0100, -0.0200, -0.0300, -0.0400, -0.0500, -0.0600, -0.0700],
                   [0.2012, -0.8897, -0.0263, -0.9401, -0.2074, -0.9463, -0.3481, -0.9343],
                   [1.7057, -0.6271, 1.2145, -1.3897, 0.6303, -1.7647, 0.1173, -1.8985],
                   [2.1731, 1.6397, 2.7358, -0.2854, 2.1840, -1.7183, 1.3018, -2.4871],
                   [-0.2717, 3.6173, 2.9206, 2.1988, 3.6638, -0.3858, 2.9155, -2.2980],
                   [-3.9859, 2.1580, 0.7984, 4.4904, 4.1181, 2.0252, 4.4782, -1.1253],
               ]       )

           tf.debugging.assert_near(query_layer[0, 0, :6, :8]		,				A__		,				atol=self.tolerance       )
           tf.debugging.assert_near(key_layer[0, 0, :6, :8]		,				A__		,				atol=self.tolerance       )



'''simple docstring'''







import tensorflow as tf

from ...tf_utils import shape_list


class 					A__      (      tf.keras.layers.Layer			):



    def __init__(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__=1		,				UpperCamelCase__=False		,				**UpperCamelCase__       )						->      Union[str, Any]:





           '''simple docstring'''
           super().__init__(**UpperCamelCase__       )

           A_						=      vocab_size
           A_						=      d_embed
           A_						=      d_proj

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

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

           A_						=      []
           A_						=      []



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





           '''simple docstring'''
           if self.n_clusters > 0:
                  A_						=      self.add_weight(
                      shape=(self.n_clusters, self.d_embed)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name="""cluster_weight"""       )
                  A_						=      self.add_weight(
                      shape=(self.n_clusters,)		,				initializer="""zeros"""		,				trainable=UpperCamelCase__		,				name="""cluster_bias"""       )

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

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



    @staticmethod
    def 					snake_case_			(					UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__=None       )						->      List[Any]:





           '''simple docstring'''
           A_						=      x
           if proj is not None:
                  A_						=      tf.einsum("""ibd,ed->ibe"""		,				UpperCamelCase__		,				UpperCamelCase__       )
           return tf.einsum("""ibd,nd->ibn"""		,				UpperCamelCase__		,				UpperCamelCase__       ) + b



    @staticmethod
    def 					snake_case_			(					UpperCamelCase__		,				UpperCamelCase__       )						->      List[str]:





           '''simple docstring'''
           A_						=      shape_list(UpperCamelCase__       )
           A_						=      tf.range(lp_size[0]		,				dtype=target.dtype       )
           A_						=      tf.stack([r, target]		,				1       )
           return tf.gather_nd(UpperCamelCase__		,				UpperCamelCase__       )







    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__		,				UpperCamelCase__=True		,				UpperCamelCase__=False       )						->      Optional[int]:





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

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

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

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

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

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

           return out
