




"""simple docstring"""





import tensorflow as tf

from ...tf_utils import shape_list



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


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


						"""simple docstring"""




						super().__init__(**UpperCamelCase_)

						__UpperCAmelCase	: int		=					vocab_size
						__UpperCAmelCase	: str		=					d_embed
						__UpperCAmelCase	: Dict		=					d_proj

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

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

						__UpperCAmelCase	: str		=					[]
						__UpperCAmelCase	: str		=					[]


			def       a_     (      self      :    int      ,  UpperCamelCase_      :    List[str]):


						"""simple docstring"""




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

												__UpperCAmelCase	: Tuple		=					self.add_weight(
												    shape=(d_emb_i, self.d_proj)      ,  initializer="zeros"      ,  trainable=UpperCamelCase_      ,  name=F"out_projs_._{i}")
												self.out_projs.append(UpperCamelCase_)
												__UpperCAmelCase	: Any		=					self.add_weight(
												    shape=(r_idx - l_idx, d_emb_i)      ,  initializer="zeros"      ,  trainable=UpperCamelCase_      ,  name=F"out_layers_._{i}_._weight"      ,  )
												__UpperCAmelCase	: Any		=					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       a_     (      UpperCamelCase_      :    int      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    Dict=None):


						"""simple docstring"""




						__UpperCAmelCase	: str		=					x
						if proj is not None:
									__UpperCAmelCase	: Optional[Any]		=					tf.einsum("ibd,ed->ibe"      ,  UpperCamelCase_      ,  UpperCamelCase_)
						return tf.einsum("ibd,nd->ibn"      ,  UpperCamelCase_      ,  UpperCamelCase_) + b


			@staticmethod
			def       a_     (      UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Optional[Any]):


						"""simple docstring"""




						__UpperCAmelCase	: int		=					shape_list(UpperCamelCase_)
						__UpperCAmelCase	: Tuple		=					tf.range(lp_size[0]      ,  dtype=target.dtype)
						__UpperCAmelCase	: List[Any]		=					tf.stack([r, target]      ,  1)
						return tf.gather_nd(UpperCamelCase_      ,  UpperCamelCase_)


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


						"""simple docstring"""




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

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

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

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

						if target is not None:
									if return_mean:
												__UpperCAmelCase	: List[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








"""simple docstring"""





def 					_UpperCamelCase    (			UpperCamelCase							)			->			int:


			"""simple docstring"""

			if not grid or not grid[0]:
						raise TypeError("The grid does not contain the appropriate information"							)

			for cell_n in range(1				,       len(grid[0]							)							):
						grid[0][cell_n] += grid[0][cell_n - 1]
			__UpperCAmelCase	: int		=					grid[0]

			for row_n in range(1				,       len(UpperCamelCase							)							):
						__UpperCAmelCase	: int		=					grid[row_n]
						__UpperCAmelCase	: str		=					fill_row(UpperCamelCase				,       UpperCamelCase							)
						__UpperCAmelCase	: List[str]		=					grid[row_n]

			return grid[-1][-1]

def 					_UpperCamelCase    (			UpperCamelCase				,       UpperCamelCase							)			->			list:


			"""simple docstring"""
			current_row[0] += row_above[0]
			for cell_n in range(1				,       len(UpperCamelCase							)							):
						current_row[cell_n] += min(current_row[cell_n - 1]				,       row_above[cell_n]							)

			return current_row


if __name__ == "__main__":
					import doctest

					doctest.testmod()


