
'''simple docstring'''






def 		A_       (	_lowerCAmelCase   :    int				,	_lowerCAmelCase   :    bool = False							):



      """simple docstring"""




      if n == 2:
            return True
      if not n % 2 or n < 2:
            return False
      if n > 5 and n % 10 not in (1, 3, 7, 9):  # can quickly check last digit
            return False
      if n > 3317044064679887385961981 and not allow_probable:
            raise ValueError(
                "Warning: upper bound of deterministic test is exceeded. "
                "Pass allow_probable=True to allow probabilistic test. "
                "A return value of True indicates a probable prime."							)
      # array bounds provided by analysis
      _lowerCamelCase				:   Any								=   [
          2047,
          1373653,
          25326001,
          3215031751,
          2152302898747,
          3474749660383,
          341550071728321,
          1,
          3825123056546413051,
          1,
          1,
          318665857834031151167461,
          3317044064679887385961981,
      ]

      _lowerCamelCase				:   Union[str, Any]								=   [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41]
      for idx, _p in enumerate(_lowerCAmelCase				,	1							):
            if n < _p:
                  # then we have our last prime to check
                  _lowerCamelCase				:   str								=   primes[:idx]
                  break
      _lowerCamelCase						,		_lowerCamelCase				:   List[str]								=   n - 1, 0
      # break up n -1 into a power of 2 (s) and
      # remaining odd component
      # essentially, solve for d * 2 ** s == n - 1
      while d % 2 == 0:
            d //= 2
            s += 1
      for prime in plist:
            _lowerCamelCase				:   int								=   False
            for r in range(_lowerCAmelCase							):
                  _lowerCamelCase				:   Any								=   pow(_lowerCAmelCase				,	d * 2**r				,	_lowerCAmelCase							)
                  # see article for analysis explanation for m
                  if (r == 0 and m == 1) or ((m + 1) % n == 0):
                        _lowerCamelCase				:   Optional[Any]								=   True
                        # this loop will not determine compositeness
                        break
            if pr:
                  continue
            # if pr is False, then the above loop never evaluated to true,
            # and the n MUST be composite
            return False
      return True







def 		A_       (	):



      """simple docstring"""




      assert not miller_rabin(561							)
      assert miller_rabin(563							)
      # 2047

      assert not miller_rabin(838201							)
      assert miller_rabin(838207							)
      # 1_373_653

      assert not miller_rabin(17316001							)
      assert miller_rabin(17316017							)
      # 25_326_001

      assert not miller_rabin(3078386641							)
      assert miller_rabin(3078386653							)
      # 3_215_031_751

      assert not miller_rabin(1713045574801							)
      assert miller_rabin(1713045574819							)
      # 2_152_302_898_747

      assert not miller_rabin(2779799728307							)
      assert miller_rabin(2779799728327							)
      # 3_474_749_660_383

      assert not miller_rabin(113850023909441							)
      assert miller_rabin(113850023909527							)
      # 341_550_071_728_321

      assert not miller_rabin(1275041018848804351							)
      assert miller_rabin(1275041018848804391							)
      # 3_825_123_056_546_413_051

      assert not miller_rabin(79666464458507787791867							)
      assert miller_rabin(79666464458507787791951							)
      # 318_665_857_834_031_151_167_461

      assert not miller_rabin(552840677446647897660333							)
      assert miller_rabin(552840677446647897660359							)
      # 3_317_044_064_679_887_385_961_981
      # upper limit for probabilistic test


if __name__ == "__main__":
 test_miller_rabin()

'''simple docstring'''






import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
 import jax.numpy as jnp

 from transformers.models.albert.modeling_flax_albert import (
     FlaxAlbertForMaskedLM,
     FlaxAlbertForMultipleChoice,
     FlaxAlbertForPreTraining,
     FlaxAlbertForQuestionAnswering,
     FlaxAlbertForSequenceClassification,
     FlaxAlbertForTokenClassification,
     FlaxAlbertModel,
 )

class  UpperCAmelCase__						(       unittest.TestCase ):



    def __init__(	self      :       Union[str, Any],__A      :       Union[str, Any],__A      :       Optional[int]=1_3,__A      :       Optional[Any]=7,__A      :       int=True,__A      :       Any=True,__A      :       Dict=True,__A      :       List[Any]=True,__A      :       Optional[Any]=9_9,__A      :       List[str]=3_2,__A      :       int=5,__A      :       str=4,__A      :       int=3_7,__A      :       Optional[Any]="gelu",__A      :       Union[str, Any]=0.1,__A      :       str=0.1,__A      :       Any=5_1_2,__A      :       Any=1_6,__A      :       Optional[int]=2,__A      :       int=0.02,__A      :       Optional[int]=4,):
          _lowerCamelCase				:   Union[str, Any]								=   parent
          _lowerCamelCase				:   Tuple								=   batch_size
          _lowerCamelCase				:   List[str]								=   seq_length
          _lowerCamelCase				:   Optional[int]								=   is_training
          _lowerCamelCase				:   str								=   use_attention_mask
          _lowerCamelCase				:   str								=   use_token_type_ids
          _lowerCamelCase				:   List[str]								=   use_labels
          _lowerCamelCase				:   Tuple								=   vocab_size
          _lowerCamelCase				:   Dict								=   hidden_size
          _lowerCamelCase				:   str								=   num_hidden_layers
          _lowerCamelCase				:   List[str]								=   num_attention_heads
          _lowerCamelCase				:   Optional[Any]								=   intermediate_size
          _lowerCamelCase				:   Optional[int]								=   hidden_act
          _lowerCamelCase				:   Any								=   hidden_dropout_prob
          _lowerCamelCase				:   List[Any]								=   attention_probs_dropout_prob
          _lowerCamelCase				:   Tuple								=   max_position_embeddings
          _lowerCamelCase				:   str								=   type_vocab_size
          _lowerCamelCase				:   str								=   type_sequence_label_size
          _lowerCamelCase				:   int								=   initializer_range
          _lowerCamelCase				:   List[Any]								=   num_choices



    def   lowerCamelCase_  (	self      :       Optional[int]     ):
          _lowerCamelCase				:   List[Any]								=   ids_tensor([self.batch_size, self.seq_length],self.vocab_size     )

          _lowerCamelCase				:   Union[str, Any]								=   None
          if self.use_attention_mask:
                _lowerCamelCase				:   Tuple								=   random_attention_mask([self.batch_size, self.seq_length]     )

          _lowerCamelCase				:   Any								=   None
          if self.use_token_type_ids:
                _lowerCamelCase				:   int								=   ids_tensor([self.batch_size, self.seq_length],self.type_vocab_size     )

          _lowerCamelCase				:   int								=   AlbertConfig(
              vocab_size=self.vocab_size,hidden_size=self.hidden_size,num_hidden_layers=self.num_hidden_layers,num_attention_heads=self.num_attention_heads,intermediate_size=self.intermediate_size,hidden_act=self.hidden_act,hidden_dropout_prob=self.hidden_dropout_prob,attention_probs_dropout_prob=self.attention_probs_dropout_prob,max_position_embeddings=self.max_position_embeddings,type_vocab_size=self.type_vocab_size,is_decoder=__A,initializer_range=self.initializer_range,)

          return config, input_ids, token_type_ids, attention_mask



    def   lowerCamelCase_  (	self      :       Dict     ):
          _lowerCamelCase				:   str								=   self.prepare_config_and_inputs()
          _lowerCamelCase						,		_lowerCamelCase						,		_lowerCamelCase						,		_lowerCamelCase				:   Union[str, Any]								=   config_and_inputs
          _lowerCamelCase				:   Any								=   {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
          return config, inputs_dict


@require_flax
class  UpperCAmelCase__						(       A						,	unittest.TestCase ):
    lowerCAmelCase_				    =		(
        (
            FlaxAlbertModel,
            FlaxAlbertForPreTraining,
            FlaxAlbertForMaskedLM,
            FlaxAlbertForMultipleChoice,
            FlaxAlbertForQuestionAnswering,
            FlaxAlbertForSequenceClassification,
            FlaxAlbertForTokenClassification,
            FlaxAlbertForQuestionAnswering,
        )
        if is_flax_available()
        else ()
    )



    def   lowerCamelCase_  (	self      :       Any     ):
          _lowerCamelCase				:   List[str]								=   FlaxAlbertModelTester(self     )



    @slow
    def   lowerCamelCase_  (	self      :       Union[str, Any]     ):
          for model_class_name in self.all_model_classes:
                _lowerCamelCase				:   str								=   model_class_name.from_pretrained("albert-base-v2"     )
                _lowerCamelCase				:   Any								=   model(np.ones((1, 1)     )     )
                self.assertIsNotNone(__A     )


@require_flax
class  UpperCAmelCase__						(       unittest.TestCase ):



    @slow
    def   lowerCamelCase_  (	self      :       Tuple     ):
          _lowerCamelCase				:   Optional[int]								=   FlaxAlbertModel.from_pretrained("albert-base-v2"     )
          _lowerCamelCase				:   int								=   np.array([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]     )
          _lowerCamelCase				:   Union[str, Any]								=   np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]     )
          _lowerCamelCase				:   Optional[Any]								=   model(__A,attention_mask=__A     )[0]
          _lowerCamelCase				:   Dict								=   (1, 1_1, 7_6_8)
          self.assertEqual(output.shape,__A     )
          _lowerCamelCase				:   Tuple								=   np.array(
              [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]]     )

          self.assertTrue(jnp.allclose(output[:, 1:4, 1:4],__A,atol=1e-4     )     )