



'''simple docstring'''
import tensorflow as tf

from ...tf_utils import shape_list






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





	def __init__(			self :    Any			,lowerCamelCase__ :    Optional[int]			,lowerCamelCase__ :    Dict			,lowerCamelCase__ :    int			,lowerCamelCase__ :    Union[str, Any]			,lowerCamelCase__ :    Optional[Any]=1			,lowerCamelCase__ :    Union[str, Any]=False			,**lowerCamelCase__ :    Union[str, Any]			):







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

							_UpperCamelCase :		Tuple			    =       vocab_size
							_UpperCamelCase :		List[str]			    =       d_embed
							_UpperCamelCase :		List[Any]			    =       d_proj

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

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

							_UpperCamelCase :		Tuple			    =       []
							_UpperCamelCase :		Optional[int]			    =       []





	def      UpperCamelCase_  (			self :    Tuple			,lowerCamelCase__ :    Optional[int]			):







							'''simple docstring'''
							if self.n_clusters > 0:
													_UpperCamelCase :		Dict			    =       self.add_weight(
													    shape=(self.n_clusters, self.d_embed)			,initializer='zeros'			,trainable=lowerCamelCase__			,name='cluster_weight'			)
													_UpperCamelCase :		List[str]			    =       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:
																									_UpperCamelCase :		str			    =       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__			)
																			_UpperCamelCase :		Optional[int]			    =       self.add_weight(
																			    shape=(self.vocab_size, self.d_embed)			,initializer='zeros'			,trainable=lowerCamelCase__			,name=F'out_layers_._{i}_._weight'			,)
																			_UpperCamelCase :		Dict			    =       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			)			):
																			_UpperCamelCase    ,      _UpperCamelCase :		int			    =       self.cutoff_ends[i], self.cutoff_ends[i + 1]
																			_UpperCamelCase :		List[Any]			    =       self.d_embed // (self.div_val**i)

																			_UpperCamelCase :		List[Any]			    =       self.add_weight(
																			    shape=(d_emb_i, self.d_proj)			,initializer='zeros'			,trainable=lowerCamelCase__			,name=F'out_projs_._{i}'			)
																			self.out_projs.append(lowerCamelCase__			)
																			_UpperCamelCase :		Union[str, Any]			    =       self.add_weight(
																			    shape=(r_idx - l_idx, d_emb_i)			,initializer='zeros'			,trainable=lowerCamelCase__			,name=F'out_layers_._{i}_._weight'			,)
																			_UpperCamelCase :		str			    =       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      UpperCamelCase_  (			lowerCamelCase__ :    Optional[Any]			,lowerCamelCase__ :    List[Any]			,lowerCamelCase__ :    Any			,lowerCamelCase__ :    Tuple=None			):







							'''simple docstring'''
							_UpperCamelCase :		List[str]			    =       x
							if proj is not None:
													_UpperCamelCase :		Optional[Any]			    =       tf.einsum('ibd,ed->ibe'			,lowerCamelCase__			,lowerCamelCase__			)
							return tf.einsum('ibd,nd->ibn'			,lowerCamelCase__			,lowerCamelCase__			) + b





	@staticmethod
	def      UpperCamelCase_  (			lowerCamelCase__ :    Dict			,lowerCamelCase__ :    Any			):







							'''simple docstring'''
							_UpperCamelCase :		Tuple			    =       shape_list(lowerCamelCase__			)
							_UpperCamelCase :		str			    =       tf.range(lp_size[0]			,dtype=target.dtype			)
							_UpperCamelCase :		Union[str, Any]			    =       tf.stack([r, target]			,1			)
							return tf.gather_nd(lowerCamelCase__			,lowerCamelCase__			)





	def      UpperCamelCase_  (			self :    Union[str, Any]			,lowerCamelCase__ :    int			,lowerCamelCase__ :    Optional[int]			,lowerCamelCase__ :    int=True			,lowerCamelCase__ :    Tuple=False			):







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

																			if self.div_val == 1:
																									_UpperCamelCase :		str			    =       self.out_layers[0][0][l_idx:r_idx]
																									_UpperCamelCase :		List[Any]			    =       self.out_layers[0][1][l_idx:r_idx]
																			else:
																									_UpperCamelCase :		List[Any]			    =       self.out_layers[i][0]
																									_UpperCamelCase :		int			    =       self.out_layers[i][1]

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

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

							if target is not None:
													if return_mean:
																			_UpperCamelCase :		List[Any]			    =       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







'''simple docstring'''
from abc import ABC, abstractmethod
from argparse import ArgumentParser






class 							lowercase__     (							lowercase       ):





	@staticmethod
	@abstractmethod
	def      UpperCamelCase_  (			lowerCamelCase__ :    ArgumentParser			):







							'''simple docstring'''
							raise NotImplementedError()





	@abstractmethod
	def      UpperCamelCase_  (			self :    Any			):







							'''simple docstring'''
							raise NotImplementedError()


