

"""simple docstring"""




import tensorflow as tf

from ...tf_utils import shape_list


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





      def __init__(  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__=1    ,						SCREAMING_SNAKE_CASE__=False    ,						**SCREAMING_SNAKE_CASE__		)   ->					Dict:
         super().__init__(**SCREAMING_SNAKE_CASE__		)

         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 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__		)   ->					Any:
         if self.n_clusters > 0:
            lowerCamelCase__						=			self.add_weight(
                shape=(self.n_clusters, self.d_embed)    ,						initializer="zeros"    ,						trainable=SCREAMING_SNAKE_CASE__    ,						name="cluster_weight"		)
            lowerCamelCase__						=			self.add_weight(
                shape=(self.n_clusters,)    ,						initializer="zeros"    ,						trainable=SCREAMING_SNAKE_CASE__    ,						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=SCREAMING_SNAKE_CASE__    ,						name=f'out_projs_._{i}'    ,						)
                  self.out_projs.append(SCREAMING_SNAKE_CASE__		)
               else:
                  self.out_projs.append(SCREAMING_SNAKE_CASE__		)
               lowerCamelCase__						=			self.add_weight(
                   shape=(self.vocab_size, self.d_embed)    ,						initializer="zeros"    ,						trainable=SCREAMING_SNAKE_CASE__    ,						name=f'out_layers_._{i}_._weight'    ,						)
               lowerCamelCase__						=			self.add_weight(
                   shape=(self.vocab_size,)    ,						initializer="zeros"    ,						trainable=SCREAMING_SNAKE_CASE__    ,						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=SCREAMING_SNAKE_CASE__    ,						name=f'out_projs_._{i}'		)
               self.out_projs.append(SCREAMING_SNAKE_CASE__		)
               lowerCamelCase__						=			self.add_weight(
                   shape=(r_idx - l_idx, d_emb_i)    ,						initializer="zeros"    ,						trainable=SCREAMING_SNAKE_CASE__    ,						name=f'out_layers_._{i}_._weight'    ,						)
               lowerCamelCase__						=			self.add_weight(
                   shape=(r_idx - l_idx,)    ,						initializer="zeros"    ,						trainable=SCREAMING_SNAKE_CASE__    ,						name=f'out_layers_._{i}_._bias'    ,						)
               self.out_layers.append((weight, bias)		)
         super().build(SCREAMING_SNAKE_CASE__		)





      @staticmethod
      def 					_lowerCamelCase  (  SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__=None		)   ->					str:
         lowerCamelCase__						=			x
         if proj is not None:
            lowerCamelCase__						=			tf.einsum("ibd,ed->ibe"    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)
         return tf.einsum("ibd,nd->ibn"    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		) + b





      @staticmethod
      def 					_lowerCamelCase  (  SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					Any:
         lowerCamelCase__						=			shape_list(SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			tf.range(lp_size[0]    ,						dtype=target.dtype		)
         lowerCamelCase__						=			tf.stack([r, target]    ,						1		)
         return tf.gather_nd(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__=True    ,						SCREAMING_SNAKE_CASE__=False		)   ->					Tuple:
         lowerCamelCase__						=			0
         if self.n_clusters == 0:
            lowerCamelCase__						=			self._logit(SCREAMING_SNAKE_CASE__    ,						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=SCREAMING_SNAKE_CASE__    ,						logits=SCREAMING_SNAKE_CASE__		)
            lowerCamelCase__						=			tf.nn.log_softmax(SCREAMING_SNAKE_CASE__    ,						axis=-1		)
         else:
            lowerCamelCase__						=			shape_list(SCREAMING_SNAKE_CASE__		)
            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(SCREAMING_SNAKE_CASE__		)
                  lowerCamelCase__						=			tf.boolean_mask(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		) - 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(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						self.out_projs[0]		)
                  lowerCamelCase__						=			tf.nn.log_softmax(SCREAMING_SNAKE_CASE__		)
                  out.append(head_logprob[..., : self.cutoffs[0]]		)
                  if target is not None:
                     lowerCamelCase__						=			tf.boolean_mask(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)
                     lowerCamelCase__						=			self._gather_logprob(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)
               else:
                  lowerCamelCase__						=			self._logit(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						self.out_projs[i]		)
                  lowerCamelCase__						=			tf.nn.log_softmax(SCREAMING_SNAKE_CASE__		)
                  lowerCamelCase__						=			self.cutoffs[0] + i - 1  # No probability for the head cluster
                  lowerCamelCase__						=			head_logprob[..., cluster_prob_idx, None] + tail_logprob
                  out.append(SCREAMING_SNAKE_CASE__		)
                  if target is not None:
                     lowerCamelCase__						=			tf.boolean_mask(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)
                     lowerCamelCase__						=			tf.boolean_mask(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)
                     lowerCamelCase__						=			self._gather_logprob(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)
                     cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
               if target is not None:
                  loss += tf.scatter_nd(SCREAMING_SNAKE_CASE__    ,						-cur_logprob    ,						shape_list(SCREAMING_SNAKE_CASE__		)		)
            lowerCamelCase__						=			tf.concat(SCREAMING_SNAKE_CASE__    ,						axis=-1		)

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

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

         return out







"""simple docstring"""




import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


SCREAMING_SNAKE_CASE_     :			Dict	  =			logging.get_logger(__name__)


class 		_A       (			__a					):
      __a								=	'encoder-decoder'
      __a								=	True





      def __init__(  self    ,						**SCREAMING_SNAKE_CASE__		)   ->					Optional[int]:
         super().__init__(**SCREAMING_SNAKE_CASE__		)
         assert (
             "encoder" in kwargs and "decoder" in kwargs
         ), "Config has to be initialized with encoder and decoder config"
         lowerCamelCase__						=			kwargs.pop("encoder"		)
         lowerCamelCase__						=			encoder_config.pop("model_type"		)
         lowerCamelCase__						=			kwargs.pop("decoder"		)
         lowerCamelCase__						=			decoder_config.pop("model_type"		)

         from ..auto.configuration_auto import AutoConfig

         lowerCamelCase__						=			AutoConfig.for_model(SCREAMING_SNAKE_CASE__    ,						**SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			AutoConfig.for_model(SCREAMING_SNAKE_CASE__    ,						**SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			True





      @classmethod
      def 					_lowerCamelCase  (  cls    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						**SCREAMING_SNAKE_CASE__		)   ->					PretrainedConfig:
         logger.info("Set `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config"		)
         lowerCamelCase__						=			True
         lowerCamelCase__						=			True

         return cls(encoder=encoder_config.to_dict()    ,						decoder=decoder_config.to_dict()    ,						**SCREAMING_SNAKE_CASE__		)





      def 					_lowerCamelCase  (  self		)   ->					int:
         lowerCamelCase__						=			copy.deepcopy(self.__dict__		)
         lowerCamelCase__						=			self.encoder.to_dict()
         lowerCamelCase__						=			self.decoder.to_dict()
         lowerCamelCase__						=			self.__class__.model_type
         return output




