

import tensorflow as tf

from ...tf_utils import shape_list

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




      """simple docstring"""

      def __init__(					self				:							Optional[Any]		,       UpperCAmelCase				:							Optional[Any]		,       UpperCAmelCase				:							Tuple		,       UpperCAmelCase				:							List[str]		,       UpperCAmelCase				:							Tuple		,       UpperCAmelCase				:							List[Any]=1		,       UpperCAmelCase				:							Optional[Any]=False		,       **UpperCAmelCase				:							Tuple     ):
            super().__init__(**__lowercase     )

            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 						__A     (					self				:							int		,       UpperCAmelCase				:							Tuple     ):






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





      @staticmethod
      def 						__A     (					UpperCAmelCase				:							Optional[int]		,       UpperCAmelCase				:							Any		,       UpperCAmelCase				:							Optional[Any]		,       UpperCAmelCase				:							int=None     ):
            A_		   =							x
            if proj is not None:
                  A_		   =							tf.einsum("ibd,ed->ibe"		,       __lowercase		,       __lowercase     )
            return tf.einsum("ibd,nd->ibn"		,       __lowercase		,       __lowercase     ) + b





      @staticmethod
      def 						__A     (					UpperCAmelCase				:							Tuple		,       UpperCAmelCase				:							Union[str, Any]     ):
            A_		   =							shape_list(__lowercase     )
            A_		   =							tf.range(lp_size[0]		,       dtype=target.dtype     )
            A_		   =							tf.stack([r, target]		,       1     )
            return tf.gather_nd(__lowercase		,       __lowercase     )







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

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

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

            return out





from torch import nn


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	if act_fn in ["swish", "silu"]:
		return nn.SiLU()
	elif act_fn == "mish":
		return nn.Mish()
	elif act_fn == "gelu":
		return nn.GELU()
	else:
		raise ValueError(f"Unsupported activation function: {act_fn}"							)
