

"""simple docstring"""


import tensorflow as tf

from ...tf_utils import shape_list





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



    '''simple docstring'''




    def __init__( self ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_=1 ,lowerCamelCase_=False ,**lowerCamelCase_    )    ->      Union[str, Any]:
          super().__init__(**lowerCamelCase_    )

          A				  =		vocab_size
          A				  =		d_embed
          A				  =		d_proj

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

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

          A				  =		[]
          A				  =		[]




    def UpperCamelCase__		( self ,lowerCamelCase_    )    ->      Union[str, Any]:



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

                      A				  =		self.add_weight(
                          shape=(d_emb_i, self.d_proj) ,initializer="""zeros""" ,trainable=lowerCamelCase_ ,name=f'out_projs_._{i}'    )
                      self.out_projs.append(lowerCamelCase_    )
                      A				  =		self.add_weight(
                          shape=(r_idx - l_idx, d_emb_i) ,initializer="""zeros""" ,trainable=lowerCamelCase_ ,name=f'out_layers_._{i}_._weight' ,)
                      A				  =		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 UpperCamelCase__		( lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_=None    )    ->      List[Any]:
          A				  =		x
          if proj is not None:
                A				  =		tf.einsum("""ibd,ed->ibe""" ,lowerCamelCase_ ,lowerCamelCase_    )
          return tf.einsum("""ibd,nd->ibn""" ,lowerCamelCase_ ,lowerCamelCase_    ) + b




    @staticmethod
    def UpperCamelCase__		( lowerCamelCase_ ,lowerCamelCase_    )    ->      List[Any]:
          A				  =		shape_list(lowerCamelCase_    )
          A				  =		tf.range(lp_size[0] ,dtype=target.dtype    )
          A				  =		tf.stack([r, target] ,1    )
          return tf.gather_nd(lowerCamelCase_ ,lowerCamelCase_    )




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

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

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

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

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




import random



def 		__A		(  lowerCamelCase_	):


		"""simple docstring"""



		SCREAMING_SNAKE_CASE :				str		=			num - 1
		SCREAMING_SNAKE_CASE :				Union[str, Any]		=			0

		while s % 2 == 0:
				SCREAMING_SNAKE_CASE :				List[str]		=			s // 2
				t += 1

		for _ in range(5	):
				SCREAMING_SNAKE_CASE :				List[Any]		=			random.randrange(2			,			num - 1	)
				SCREAMING_SNAKE_CASE :				List[Any]		=			pow(lowerCamelCase_			,			lowerCamelCase_			,			lowerCamelCase_	)
				if v != 1:
						SCREAMING_SNAKE_CASE :				Tuple		=			0
						while v != (num - 1):
								if i == t - 1:
										return False
								else:
										SCREAMING_SNAKE_CASE :				Any		=			i + 1
										SCREAMING_SNAKE_CASE :				str		=			(v**2) % num
		return True



def 		__A		(  lowerCamelCase_	):


		"""simple docstring"""



		if num < 2:
				return False

		SCREAMING_SNAKE_CASE :				List[str]		=			[
		    2,
		    3,
		    5,
		    7,
		    11,
		    13,
		    17,
		    19,
		    23,
		    29,
		    31,
		    37,
		    41,
		    43,
		    47,
		    53,
		    59,
		    61,
		    67,
		    71,
		    73,
		    79,
		    83,
		    89,
		    97,
		    1_01,
		    1_03,
		    1_07,
		    1_09,
		    1_13,
		    1_27,
		    1_31,
		    1_37,
		    1_39,
		    1_49,
		    1_51,
		    1_57,
		    1_63,
		    1_67,
		    1_73,
		    1_79,
		    1_81,
		    1_91,
		    1_93,
		    1_97,
		    1_99,
		    2_11,
		    2_23,
		    2_27,
		    2_29,
		    2_33,
		    2_39,
		    2_41,
		    2_51,
		    2_57,
		    2_63,
		    2_69,
		    2_71,
		    2_77,
		    2_81,
		    2_83,
		    2_93,
		    3_07,
		    3_11,
		    3_13,
		    3_17,
		    3_31,
		    3_37,
		    3_47,
		    3_49,
		    3_53,
		    3_59,
		    3_67,
		    3_73,
		    3_79,
		    3_83,
		    3_89,
		    3_97,
		    4_01,
		    4_09,
		    4_19,
		    4_21,
		    4_31,
		    4_33,
		    4_39,
		    4_43,
		    4_49,
		    4_57,
		    4_61,
		    4_63,
		    4_67,
		    4_79,
		    4_87,
		    4_91,
		    4_99,
		    5_03,
		    5_09,
		    5_21,
		    5_23,
		    5_41,
		    5_47,
		    5_57,
		    5_63,
		    5_69,
		    5_71,
		    5_77,
		    5_87,
		    5_93,
		    5_99,
		    6_01,
		    6_07,
		    6_13,
		    6_17,
		    6_19,
		    6_31,
		    6_41,
		    6_43,
		    6_47,
		    6_53,
		    6_59,
		    6_61,
		    6_73,
		    6_77,
		    6_83,
		    6_91,
		    7_01,
		    7_09,
		    7_19,
		    7_27,
		    7_33,
		    7_39,
		    7_43,
		    7_51,
		    7_57,
		    7_61,
		    7_69,
		    7_73,
		    7_87,
		    7_97,
		    8_09,
		    8_11,
		    8_21,
		    8_23,
		    8_27,
		    8_29,
		    8_39,
		    8_53,
		    8_57,
		    8_59,
		    8_63,
		    8_77,
		    8_81,
		    8_83,
		    8_87,
		    9_07,
		    9_11,
		    9_19,
		    9_29,
		    9_37,
		    9_41,
		    9_47,
		    9_53,
		    9_67,
		    9_71,
		    9_77,
		    9_83,
		    9_91,
		    9_97,
		]

		if num in low_primes:
				return True

		for prime in low_primes:
				if (num % prime) == 0:
						return False

		return rabin_miller(lowerCamelCase_	)



def 		__A		(  lowerCamelCase_ = 10_24	):


		"""simple docstring"""



		while True:
				SCREAMING_SNAKE_CASE :				Optional[Any]		=			random.randrange(2 ** (keysize - 1)			,			2 ** (keysize)	)
				if is_prime_low_num(lowerCamelCase_	):
						return num


if __name__ == "__main__":
							__UpperCAmelCase            =       generate_large_prime()
							print(("""Prime number:""", num))
							print(("""is_prime_low_num:""", is_prime_low_num(num)))



