


import tensorflow as tf

from ...tf_utils import shape_list
class 	lowerCamelCase_       (				tf.keras.layers.Layer				):



					'''simple docstring'''







					def __init__(				self ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_=1 ,       snake_case_=False ,       **snake_case_						)   ->			List[str]:







								'''simple docstring'''

								super().__init__(**_A						)

								__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 ,       snake_case_						)   ->			Any:







								'''simple docstring'''

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






					@staticmethod
					def A      (				snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_=None						)   ->			Union[str, Any]:







								'''simple docstring'''

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






					@staticmethod
					def A      (				snake_case_ ,       snake_case_						)   ->			Optional[Any]:







								'''simple docstring'''

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






					def A      (				self ,       snake_case_ ,       snake_case_ ,       snake_case_=True ,       snake_case_=False						)   ->			int:







								'''simple docstring'''

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

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

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

								return out






from ..utils import DummyObject, requires_backends
class 	lowerCamelCase_       (				metaclass=lowerCAmelCase__				):



					'''simple docstring'''







					__UpperCAmelCase            =				["speech"]






					def __init__(				self ,       *snake_case_ ,       **snake_case_						)   ->			List[str]:







								'''simple docstring'''

								requires_backends(self ,       ['''speech''']						)




class 	lowerCamelCase_       (				metaclass=lowerCAmelCase__				):



					'''simple docstring'''







					__UpperCAmelCase            =				["speech"]






					def __init__(				self ,       *snake_case_ ,       **snake_case_						)   ->			Tuple:







								'''simple docstring'''

								requires_backends(self ,       ['''speech''']						)


