
"""simple docstring"""



import tensorflow as tf

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






				'''simple docstring'''



				def __init__(					self	:			Optional[Any]				,     UpperCAmelCase_	:			Tuple				,     UpperCAmelCase_	:			Optional[Any]				,     UpperCAmelCase_	:			Optional[int]				,     UpperCAmelCase_	:			Any				,     UpperCAmelCase_	:			Optional[int]=1				,     UpperCAmelCase_	:			List[str]=False				,     **UpperCAmelCase_	:			Tuple  )		->  str:







											"""simple docstring"""




											super().__init__(**UpperCAmelCase_  )

											_lowerCAmelCase  	=			vocab_size
											_lowerCAmelCase  	=			d_embed
											_lowerCAmelCase  	=			d_proj

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

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

											_lowerCAmelCase  	=			[]
											_lowerCAmelCase  	=			[]


				def 				__lowerCamelCase	(					self	:			str				,     UpperCAmelCase_	:			Optional[Any]  )		->  int:







											"""simple docstring"""




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

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







											"""simple docstring"""




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


				@staticmethod
				def 				__lowerCamelCase	(					UpperCAmelCase_	:			Dict				,     UpperCAmelCase_	:			Tuple  )		->  Optional[int]:







											"""simple docstring"""




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







				def 				__lowerCamelCase	(					self	:			Union[str, Any]				,     UpperCAmelCase_	:			str				,     UpperCAmelCase_	:			Any				,     UpperCAmelCase_	:			Union[str, Any]=True				,     UpperCAmelCase_	:			Optional[Any]=False  )		->  List[str]:







											"""simple docstring"""




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

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

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

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

											if target is not None:
																		if return_mean:
																									_lowerCAmelCase  	=			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 	__snake_case  (    SCREAMING_SNAKE_CASE:					str						):






							"""simple docstring"""



							_lowerCAmelCase  	=			[0 for i in range(len(SCREAMING_SNAKE_CASE						)						)]

							# initialize interval's left pointer and right pointer
							_lowerCAmelCase	,		_lowerCAmelCase  	=			0, 0

							for i in range(1     ,  len(SCREAMING_SNAKE_CASE						)						):
														# case when current index is inside the interval
														if i <= right_pointer:
																					_lowerCAmelCase  	=			min(right_pointer - i + 1     ,  z_result[i - left_pointer]						)
																					_lowerCAmelCase  	=			min_edge

														while go_next(SCREAMING_SNAKE_CASE     ,  SCREAMING_SNAKE_CASE     ,  SCREAMING_SNAKE_CASE						):
																					z_result[i] += 1

														# if new index's result gives us more right interval,
														# we've to update left_pointer and right_pointer
														if i + z_result[i] - 1 > right_pointer:
																					_lowerCAmelCase	,		_lowerCAmelCase  	=			i, i + z_result[i] - 1

							return z_result






def 	__snake_case  (    SCREAMING_SNAKE_CASE:					int     ,  SCREAMING_SNAKE_CASE:					list[int]     ,  SCREAMING_SNAKE_CASE:					str						):






							"""simple docstring"""



							return i + z_result[i] < len(SCREAMING_SNAKE_CASE						) and s[z_result[i]] == s[i + z_result[i]]






def 	__snake_case  (    SCREAMING_SNAKE_CASE:					str     ,  SCREAMING_SNAKE_CASE:					str						):






							"""simple docstring"""



							_lowerCAmelCase  	=			0
							# concatenate 'pattern' and 'input_str' and call z_function
							# with concatenated string
							_lowerCAmelCase  	=			z_function(pattern + input_str						)

							for val in z_result:
														# if value is greater then length of the pattern string
														# that means this index is starting position of substring
														# which is equal to pattern string
														if val >= len(SCREAMING_SNAKE_CASE						):
																					answer += 1

							return answer


if __name__ == "__main__":
	import doctest

	doctest.testmod()






