





import tensorflow as tf

from ...tf_utils import shape_list

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





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase=1	,  _lowerCAmelCase=False	,  **_lowerCAmelCase	):
										super().__init__(**_lowerCAmelCase	)

										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 					__magic_name__ ( self	,  _lowerCAmelCase	):
										if self.n_clusters > 0:
															lowerCamelCase__       =       self.add_weight(
															    shape=(self.n_clusters, self.d_embed)	,  initializer="zeros"	,  trainable=_lowerCAmelCase	,  name="cluster_weight"	)
															lowerCamelCase__       =       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:
																									lowerCamelCase__       =       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	)
																				lowerCamelCase__       =       self.add_weight(
																				    shape=(self.vocab_size, self.d_embed)	,  initializer="zeros"	,  trainable=_lowerCAmelCase	,  name=F"out_layers_._{i}_._weight"	,  )
																				lowerCamelCase__       =       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	)	):
																				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=_lowerCAmelCase	,  name=F"out_projs_._{i}"	)
																				self.out_projs.append(_lowerCAmelCase	)
																				lowerCamelCase__       =       self.add_weight(
																				    shape=(r_idx - l_idx, d_emb_i)	,  initializer="zeros"	,  trainable=_lowerCAmelCase	,  name=F"out_layers_._{i}_._weight"	,  )
																				lowerCamelCase__       =       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 					__magic_name__ ( _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase=None	):
										lowerCamelCase__       =       x
										if proj is not None:
															lowerCamelCase__       =       tf.einsum("ibd,ed->ibe"	,  _lowerCAmelCase	,  _lowerCAmelCase	)
										return tf.einsum("ibd,nd->ibn"	,  _lowerCAmelCase	,  _lowerCAmelCase	) + b





					@staticmethod
					def 					__magic_name__ ( _lowerCAmelCase	,  _lowerCAmelCase	):
										lowerCamelCase__       =       shape_list(_lowerCAmelCase	)
										lowerCamelCase__       =       tf.range(lp_size[0]	,  dtype=target.dtype	)
										lowerCamelCase__       =       tf.stack([r, target]	,  1	)
										return tf.gather_nd(_lowerCAmelCase	,  _lowerCAmelCase	)






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

										if target is not None:
															if return_mean:
																				lowerCamelCase__       =       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








def    __UpperCamelCase      (    a) ->float:
					return 10 - x * x





def    __UpperCamelCase      (    a,     a) ->float:
					# Bolzano theory in order to find if there is a root between a and b
					if equation(a) * equation(a) >= 0:
										raise ValueError("Wrong space!")

					lowerCamelCase__       =       a
					while (b - a) >= 0.0_1:
										# Find middle point
										lowerCamelCase__       =       (a + b) / 2
										# Check if middle point is root
										if equation(a) == 0.0:
															break
										# Decide the side to repeat the steps
										if equation(a) * equation(a) < 0:
															lowerCamelCase__       =       c
										else:
															lowerCamelCase__       =       c
					return c


if __name__ == "__main__":
						import doctest

						doctest.testmod()

						print(bisection(-2, 5))
						print(bisection(0, 6))

