


import tensorflow as tf

from ...tf_utils import shape_list






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


   '''simple docstring'''







   def __init__(     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_=1   ,							UpperCamelCase_=False   ,							**UpperCamelCase_		):
        super().__init__(**UpperCamelCase_		)

        lowercase_       :str     =       vocab_size
        lowercase_       :Optional[int]     =       d_embed
        lowercase_       :List[str]     =       d_proj

        lowercase_       :Optional[Any]     =       cutoffs + [vocab_size]
        lowercase_       :Optional[int]     =       [0] + self.cutoffs
        lowercase_       :int     =       div_val

        lowercase_       :str     =       self.cutoffs[0]
        lowercase_       :Union[str, Any]     =       len(self.cutoffs		) - 1
        lowercase_       :Dict     =       self.shortlist_size + self.n_clusters
        lowercase_       :List[Any]     =       keep_order

        lowercase_       :Dict     =       []
        lowercase_       :Optional[int]     =       []
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_		):
        if self.n_clusters > 0:
             lowercase_       :Union[str, Any]     =       self.add_weight(
                 shape=(self.n_clusters, self.d_embed)   ,							initializer='''zeros'''   ,							trainable=UpperCamelCase_   ,							name='''cluster_weight'''		)
             lowercase_       :Dict     =       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:
                       lowercase_       :Dict     =       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_		)
                  lowercase_       :List[str]     =       self.add_weight(
                      shape=(self.vocab_size, self.d_embed)   ,							initializer='''zeros'''   ,							trainable=UpperCamelCase_   ,							name=f"out_layers_._{i}_._weight"   ,							)
                  lowercase_       :List[str]     =       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		)		):
                  lowercase_       :int     =       self.cutoff_ends[i], self.cutoff_ends[i + 1]
                  lowercase_       :Tuple     =       self.d_embed // (self.div_val**i)

                  lowercase_       :Union[str, Any]     =       self.add_weight(
                      shape=(d_emb_i, self.d_proj)   ,							initializer='''zeros'''   ,							trainable=UpperCamelCase_   ,							name=f"out_projs_._{i}"		)
                  self.out_projs.append(UpperCamelCase_		)
                  lowercase_       :Dict     =       self.add_weight(
                      shape=(r_idx - l_idx, d_emb_i)   ,							initializer='''zeros'''   ,							trainable=UpperCamelCase_   ,							name=f"out_layers_._{i}_._weight"   ,							)
                  lowercase_       :Dict     =       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 		UpperCamelCase  (     UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_=None		):
        lowercase_       :Optional[int]     =       x
        if proj is not None:
             lowercase_       :str     =       tf.einsum('''ibd,ed->ibe'''   ,							UpperCamelCase_   ,							UpperCamelCase_		)
        return tf.einsum('''ibd,nd->ibn'''   ,							UpperCamelCase_   ,							UpperCamelCase_		) + b
   @staticmethod
   def 		UpperCamelCase  (     UpperCamelCase_   ,							UpperCamelCase_		):
        lowercase_       :List[str]     =       shape_list(UpperCamelCase_		)
        lowercase_       :Optional[int]     =       tf.range(lp_size[0]   ,							dtype=target.dtype		)
        lowercase_       :str     =       tf.stack([r, target]   ,							1		)
        return tf.gather_nd(UpperCamelCase_   ,							UpperCamelCase_		)




   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_=True   ,							UpperCamelCase_=False		):
        lowercase_       :Optional[Any]     =       0
        if self.n_clusters == 0:
             lowercase_       :Any     =       self._logit(UpperCamelCase_   ,							self.out_layers[0][0]   ,							self.out_layers[0][1]   ,							self.out_projs[0]		)
             if target is not None:
                  lowercase_       :Optional[int]     =       tf.nn.sparse_softmax_cross_entropy_with_logits(labels=UpperCamelCase_   ,							logits=UpperCamelCase_		)
             lowercase_       :str     =       tf.nn.log_softmax(UpperCamelCase_   ,							axis=-1		)
        else:
             lowercase_       :Tuple     =       shape_list(UpperCamelCase_		)
             lowercase_       :Tuple     =       []
             lowercase_       :List[str]     =       tf.zeros(hidden_sizes[:2]		)
             for i in range(len(self.cutoffs		)		):
                  lowercase_       :Optional[Any]     =       self.cutoff_ends[i], self.cutoff_ends[i + 1]
                  if target is not None:
                       lowercase_       :str     =       (target >= l_idx) & (target < r_idx)
                       lowercase_       :List[str]     =       tf.where(UpperCamelCase_		)
                       lowercase_       :int     =       tf.boolean_mask(UpperCamelCase_   ,							UpperCamelCase_		) - l_idx

                  if self.div_val == 1:
                       lowercase_       :Optional[Any]     =       self.out_layers[0][0][l_idx:r_idx]
                       lowercase_       :Dict     =       self.out_layers[0][1][l_idx:r_idx]
                  else:
                       lowercase_       :int     =       self.out_layers[i][0]
                       lowercase_       :List[Any]     =       self.out_layers[i][1]

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

                       lowercase_       :Optional[Any]     =       self._logit(UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							self.out_projs[0]		)
                       lowercase_       :Any     =       tf.nn.log_softmax(UpperCamelCase_		)
                       out.append(head_logprob[..., : self.cutoffs[0]]		)
                       if target is not None:
                            lowercase_       :Tuple     =       tf.boolean_mask(UpperCamelCase_   ,							UpperCamelCase_		)
                            lowercase_       :int     =       self._gather_logprob(UpperCamelCase_   ,							UpperCamelCase_		)
                  else:
                       lowercase_       :int     =       self._logit(UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							self.out_projs[i]		)
                       lowercase_       :Optional[Any]     =       tf.nn.log_softmax(UpperCamelCase_		)
                       lowercase_       :Dict     =       self.cutoffs[0] + i - 1  # No probability for the head cluster
                       lowercase_       :Optional[int]     =       head_logprob[..., cluster_prob_idx, None] + tail_logprob
                       out.append(UpperCamelCase_		)
                       if target is not None:
                            lowercase_       :Dict     =       tf.boolean_mask(UpperCamelCase_   ,							UpperCamelCase_		)
                            lowercase_       :Optional[Any]     =       tf.boolean_mask(UpperCamelCase_   ,							UpperCamelCase_		)
                            lowercase_       :str     =       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_		)		)
             lowercase_       :Any     =       tf.concat(UpperCamelCase_   ,							axis=-1		)

        if target is not None:
             if return_mean:
                  lowercase_       :Any     =       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





import unittest

import numpy as np

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
   from transformers.models.bert.modeling_flax_bert import (
       FlaxBertForMaskedLM,
       FlaxBertForMultipleChoice,
       FlaxBertForNextSentencePrediction,
       FlaxBertForPreTraining,
       FlaxBertForQuestionAnswering,
       FlaxBertForSequenceClassification,
       FlaxBertForTokenClassification,
       FlaxBertModel,
   )






class      UpperCamelCase					(		unittest.TestCase     ):


   '''simple docstring'''







   def __init__(     self   ,							UpperCamelCase_   ,							UpperCamelCase_=13   ,							UpperCamelCase_=7   ,							UpperCamelCase_=True   ,							UpperCamelCase_=True   ,							UpperCamelCase_=True   ,							UpperCamelCase_=True   ,							UpperCamelCase_=99   ,							UpperCamelCase_=32   ,							UpperCamelCase_=5   ,							UpperCamelCase_=4   ,							UpperCamelCase_=37   ,							UpperCamelCase_="gelu"   ,							UpperCamelCase_=0.1   ,							UpperCamelCase_=0.1   ,							UpperCamelCase_=512   ,							UpperCamelCase_=16   ,							UpperCamelCase_=2   ,							UpperCamelCase_=0.02   ,							UpperCamelCase_=4   ,							):
        lowercase_       :Union[str, Any]     =       parent
        lowercase_       :int     =       batch_size
        lowercase_       :int     =       seq_length
        lowercase_       :str     =       is_training
        lowercase_       :Dict     =       use_attention_mask
        lowercase_       :List[Any]     =       use_token_type_ids
        lowercase_       :str     =       use_labels
        lowercase_       :str     =       vocab_size
        lowercase_       :Optional[int]     =       hidden_size
        lowercase_       :Dict     =       num_hidden_layers
        lowercase_       :List[str]     =       num_attention_heads
        lowercase_       :int     =       intermediate_size
        lowercase_       :Union[str, Any]     =       hidden_act
        lowercase_       :Optional[int]     =       hidden_dropout_prob
        lowercase_       :Tuple     =       attention_probs_dropout_prob
        lowercase_       :int     =       max_position_embeddings
        lowercase_       :List[Any]     =       type_vocab_size
        lowercase_       :Any     =       type_sequence_label_size
        lowercase_       :Tuple     =       initializer_range
        lowercase_       :Any     =       num_choices
   def 		UpperCamelCase  (     self		):
        lowercase_       :Optional[int]     =       ids_tensor([self.batch_size, self.seq_length]   ,							self.vocab_size		)

        lowercase_       :List[str]     =       None
        if self.use_attention_mask:
             lowercase_       :Any     =       random_attention_mask([self.batch_size, self.seq_length]		)

        lowercase_       :List[Any]     =       None
        if self.use_token_type_ids:
             lowercase_       :Union[str, Any]     =       ids_tensor([self.batch_size, self.seq_length]   ,							self.type_vocab_size		)

        lowercase_       :Union[str, Any]     =       BertConfig(
            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=UpperCamelCase_   ,							initializer_range=self.initializer_range   ,							)

        return config, input_ids, token_type_ids, attention_mask
   def 		UpperCamelCase  (     self		):
        lowercase_       :Optional[Any]     =       self.prepare_config_and_inputs()
        lowercase_    ,      lowercase_    ,      lowercase_    ,      lowercase_       :List[str]     =       config_and_inputs
        lowercase_       :Union[str, Any]     =       {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
        return config, inputs_dict




   def 		UpperCamelCase  (     self		):
        lowercase_       :Dict     =       self.prepare_config_and_inputs()
        lowercase_    ,      lowercase_    ,      lowercase_    ,      lowercase_       :Dict     =       config_and_inputs

        lowercase_       :str     =       True
        lowercase_       :Dict     =       floats_tensor([self.batch_size, self.seq_length, self.hidden_size]		)
        lowercase_       :Optional[Any]     =       ids_tensor([self.batch_size, self.seq_length]   ,							vocab_size=2		)

        return (
            config,
            input_ids,
            attention_mask,
            encoder_hidden_states,
            encoder_attention_mask,
        )






@require_flax
class      UpperCamelCase					(		lowercase__						,      unittest.TestCase     ):


   '''simple docstring'''







   lowercase					:	int									=True

   lowercase					:	Dict									=(
       (
           FlaxBertModel,
           FlaxBertForPreTraining,
           FlaxBertForMaskedLM,
           FlaxBertForMultipleChoice,
           FlaxBertForQuestionAnswering,
           FlaxBertForNextSentencePrediction,
           FlaxBertForSequenceClassification,
           FlaxBertForTokenClassification,
           FlaxBertForQuestionAnswering,
       )
       if is_flax_available()
       else ()
   )
   def 		UpperCamelCase  (     self		):
        lowercase_       :Union[str, Any]     =       FlaxBertModelTester(self		)




   @slow
   def 		UpperCamelCase  (     self		):
        # Only check this for base model, not necessary for all model classes.
        # This will also help speed-up tests.
        lowercase_       :Dict     =       FlaxBertModel.from_pretrained('''bert-base-cased'''		)
        lowercase_       :str     =       model(np.ones((1, 1)		)		)
        self.assertIsNotNone(UpperCamelCase_		)

