



'''simple docstring'''


import tensorflow as tf

from ...tf_utils import shape_list




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







	def __init__(    self ,				_lowerCAmelCase ,				_lowerCAmelCase ,				_lowerCAmelCase ,				_lowerCAmelCase ,				_lowerCAmelCase=1 ,				_lowerCAmelCase=False ,				**_lowerCAmelCase						)    ->    Tuple:
							'''simple docstring'''






							super().__init__(**_lowerCAmelCase						)

							lowercase        =							vocab_size
							lowercase        =							d_embed
							lowercase        =							d_proj

							lowercase        =							cutoffs + [vocab_size]
							lowercase        =							[0] + self.cutoffs
							lowercase        =							div_val

							lowercase        =							self.cutoffs[0]
							lowercase        =							len(self.cutoffs						) - 1
							lowercase        =							self.shortlist_size + self.n_clusters
							lowercase        =							keep_order

							lowercase        =							[]
							lowercase        =							[]







	def   _a   (    self ,				_lowerCAmelCase						)    ->    str:
							'''simple docstring'''






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

																			lowercase        =							self.add_weight(
																			    shape=(d_emb_i, self.d_proj) ,				initializer="""zeros""" ,				trainable=_lowerCAmelCase ,				name=F"""out_projs_._{i}"""						)
																			self.out_projs.append(_lowerCAmelCase						)
																			lowercase        =							self.add_weight(
																			    shape=(r_idx - l_idx, d_emb_i) ,				initializer="""zeros""" ,				trainable=_lowerCAmelCase ,				name=F"""out_layers_._{i}_._weight""" ,				)
																			lowercase        =							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   _a   (    _lowerCAmelCase ,				_lowerCAmelCase ,				_lowerCAmelCase ,				_lowerCAmelCase=None						)    ->    Optional[int]:
							'''simple docstring'''






							lowercase        =							x
							if proj is not None:
													lowercase        =							tf.einsum("""ibd,ed->ibe""" ,				_lowerCAmelCase ,				_lowerCAmelCase						)
							return tf.einsum("""ibd,nd->ibn""" ,				_lowerCAmelCase ,				_lowerCAmelCase						) + b







	@staticmethod
	def   _a   (    _lowerCAmelCase ,				_lowerCAmelCase						)    ->    Any:
							'''simple docstring'''






							lowercase        =							shape_list(_lowerCAmelCase						)
							lowercase        =							tf.range(lp_size[0] ,				dtype=target.dtype						)
							lowercase        =							tf.stack([r, target] ,				1						)
							return tf.gather_nd(_lowerCAmelCase ,				_lowerCAmelCase						)







	def   _a   (    self ,				_lowerCAmelCase ,				_lowerCAmelCase ,				_lowerCAmelCase=True ,				_lowerCAmelCase=False						)    ->    Dict:
							'''simple docstring'''






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

																			if self.div_val == 1:
																									lowercase        =							self.out_layers[0][0][l_idx:r_idx]
																									lowercase        =							self.out_layers[0][1][l_idx:r_idx]
																			else:
																									lowercase        =							self.out_layers[i][0]
																									lowercase        =							self.out_layers[i][1]

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

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

							if target is not None:
													if return_mean:
																			lowercase        =							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 transformers import HfArgumentParser, TensorFlowBenchmark, TensorFlowBenchmarkArguments







def      SCREAMING_SNAKE_CASE				(				):
						lowercase        =							HfArgumentParser(lowercase_						)
						lowercase        =							parser.parse_args_into_dataclasses()[0]
						lowercase        =							TensorFlowBenchmark(args=lowercase_						)
						try:
												lowercase        =							parser.parse_args_into_dataclasses()[0]
						except ValueError as e:
												lowercase        =							"""Arg --no_{0} is no longer used, please use --no-{0} instead."""
												lowercase        =							""" """.join(str(lowercase_						).split(""" """						)[:-1]						)
												lowercase        =							""""""
												lowercase        =							eval(str(lowercase_						).split(""" """						)[-1]						)
												lowercase        =							[]
												for arg in depreciated_args:
																		# arg[2:] removes '--'
																		if arg[2:] in TensorFlowBenchmark.deprecated_args:
																								# arg[5:] removes '--no_'
																								full_error_msg += arg_error_msg.format(arg[5:]						)
																		else:
																								wrong_args.append(lowercase_						)
												if len(lowercase_						) > 0:
																		lowercase        =							full_error_msg + begin_error_msg + str(lowercase_						)
												raise ValueError(lowercase_						)
						benchmark.run()


if __name__ == "__main__":
						main()


