

"""simple docstring"""





def   __SCREAMING_SNAKE_CASE  (      A_						,     A_ = False ):
	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 > 3_31_70_44_06_46_79_88_73_85_96_19_81 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__	:						Optional[int]  				=	[
	    20_47,
	    1_37_36_53,
	    25_32_60_01,
	    32_15_03_17_51,
	    2_15_23_02_89_87_47,
	    3_47_47_49_66_03_83,
	    3_41_55_00_71_72_83_21,
	    1,
	    3_82_51_23_05_65_46_41_30_51,
	    1,
	    1,
	    31_86_65_85_78_34_03_11_51_16_74_61,
	    3_31_70_44_06_46_79_88_73_85_96_19_81,
	]

	lowerCAmelCase__	:						int  				=	[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__	:						List[str]  				=	primes[:idx]
			break
	lowerCAmelCase__	:						Optional[Any]  				=	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__	:						str  				=	False
		for r in range(_lowerCamelCase ):
			lowerCAmelCase__	:						int  				=	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__	:						Tuple  				=	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   __SCREAMING_SNAKE_CASE  (      ):
	assert not miller_rabin(5_61 )
	assert miller_rabin(5_63 )
	# 2047

	assert not miller_rabin(83_82_01 )
	assert miller_rabin(83_82_07 )
	# 1_373_653

	assert not miller_rabin(17_31_60_01 )
	assert miller_rabin(17_31_60_17 )
	# 25_326_001

	assert not miller_rabin(30_78_38_66_41 )
	assert miller_rabin(30_78_38_66_53 )
	# 3_215_031_751

	assert not miller_rabin(1_71_30_45_57_48_01 )
	assert miller_rabin(1_71_30_45_57_48_19 )
	# 2_152_302_898_747

	assert not miller_rabin(2_77_97_99_72_83_07 )
	assert miller_rabin(2_77_97_99_72_83_27 )
	# 3_474_749_660_383

	assert not miller_rabin(1_13_85_00_23_90_94_41 )
	assert miller_rabin(1_13_85_00_23_90_95_27 )
	# 341_550_071_728_321

	assert not miller_rabin(1_27_50_41_01_88_48_80_43_51 )
	assert miller_rabin(1_27_50_41_01_88_48_80_43_91 )
	# 3_825_123_056_546_413_051

	assert not miller_rabin(7_96_66_46_44_58_50_77_87_79_18_67 )
	assert miller_rabin(7_96_66_46_44_58_50_77_87_79_19_51 )
	# 318_665_857_834_031_151_167_461

	assert not miller_rabin(55_28_40_67_74_46_64_78_97_66_03_33 )
	assert miller_rabin(55_28_40_67_74_46_64_78_97_66_03_59 )
	# 3_317_044_064_679_887_385_961_981
	# upper limit for probabilistic test


if __name__ == "__main__":
							test_miller_rabin()





import unittest

from transformers import BigBirdConfig, 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

      from transformers.models.big_bird.modeling_flax_big_bird import (
          FlaxBigBirdForCausalLM,
          FlaxBigBirdForMaskedLM,
          FlaxBigBirdForMultipleChoice,
          FlaxBigBirdForPreTraining,
          FlaxBigBirdForQuestionAnswering,
          FlaxBigBirdForSequenceClassification,
          FlaxBigBirdForTokenClassification,
          FlaxBigBirdModel,
      )



class   snake_case_       (		unittest.TestCase       ):






   def __init__(      self				:     Tuple	,						lowercase_				:     List[Any]	,						lowercase_				:     Union[str, Any]=2	,						lowercase_				:     Union[str, Any]=56	,						lowercase_				:     Tuple=True	,						lowercase_				:     Optional[Any]=True	,						lowercase_				:     Optional[Any]=True	,						lowercase_				:     int=True	,						lowercase_				:     Any=99	,						lowercase_				:     int=32	,						lowercase_				:     str=2	,						lowercase_				:     Union[str, Any]=2	,						lowercase_				:     Dict=7	,						lowercase_				:     Dict="gelu_new"	,						lowercase_				:     Tuple=0.1	,						lowercase_				:     List[Any]=0.1	,						lowercase_				:     Tuple=5_12	,						lowercase_				:     Optional[Any]=16	,						lowercase_				:     List[Any]=2	,						lowercase_				:     Dict=0.02	,						lowercase_				:     int=4	,						lowercase_				:     Tuple="block_sparse"	,						lowercase_				:     Dict=True	,						lowercase_				:     Optional[int]=False	,						lowercase_				:     Dict=2	,						lowercase_				:     int=3	,						)					->						Union[str, Any]:
       lowercase__ :   Dict    					=				parent
       lowercase__ :   Dict    					=				batch_size
       lowercase__ :   Tuple    					=				seq_length
       lowercase__ :   Dict    					=				is_training
       lowercase__ :   Dict    					=				use_attention_mask
       lowercase__ :   Tuple    					=				use_token_type_ids
       lowercase__ :   Optional[int]    					=				use_labels
       lowercase__ :   List[Any]    					=				vocab_size
       lowercase__ :   Any    					=				hidden_size
       lowercase__ :   List[Any]    					=				num_hidden_layers
       lowercase__ :   Union[str, Any]    					=				num_attention_heads
       lowercase__ :   str    					=				intermediate_size
       lowercase__ :   int    					=				hidden_act
       lowercase__ :   str    					=				hidden_dropout_prob
       lowercase__ :   List[str]    					=				attention_probs_dropout_prob
       lowercase__ :   Optional[Any]    					=				max_position_embeddings
       lowercase__ :   Union[str, Any]    					=				type_vocab_size
       lowercase__ :   Dict    					=				type_sequence_label_size
       lowercase__ :   Any    					=				initializer_range
       lowercase__ :   List[str]    					=				num_choices

       lowercase__ :   str    					=				rescale_embeddings
       lowercase__ :   Optional[Any]    					=				attention_type
       lowercase__ :   Optional[int]    					=				use_bias
       lowercase__ :   Optional[int]    					=				block_size
       lowercase__ :   str    					=				num_random_blocks






   def       __UpperCamelCase     (      self				:     str  )					->						Optional[Any]:
       lowercase__ :   List[Any]    					=				ids_tensor([self.batch_size, self.seq_length]	,						self.vocab_size  )

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

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

       lowercase__ :   int    					=				BigBirdConfig(
           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=lowercase_	,						initializer_range=self.initializer_range	,						attention_type=self.attention_type	,						block_size=self.block_size	,						num_random_blocks=self.num_random_blocks	,						use_bias=self.use_bias	,						rescale_embeddings=self.rescale_embeddings	,						)

       return config, input_ids, token_type_ids, attention_mask







   def       __UpperCamelCase     (      self				:     Union[str, Any]  )					->						int:
       lowercase__ :   int    					=				self.prepare_config_and_inputs()
       lowercase__					,   lowercase__					,   lowercase__					,   lowercase__ :   Dict    					=				config_and_inputs
       lowercase__ :   Union[str, Any]    					=				{
           "input_ids": input_ids,
           "token_type_ids": token_type_ids,
           "attention_mask": attention_mask,
       }
       return config, inputs_dict



@require_flax
class   snake_case_       (		__A  ,unittest.TestCase       ):
   __A	:		Optional[int]         =				(
       (
           FlaxBigBirdForCausalLM,
           FlaxBigBirdModel,
           FlaxBigBirdForPreTraining,
           FlaxBigBirdForMaskedLM,
           FlaxBigBirdForMultipleChoice,
           FlaxBigBirdForQuestionAnswering,
           FlaxBigBirdForSequenceClassification,
           FlaxBigBirdForTokenClassification,
       )
       if is_flax_available()
       else ()
   )

   __A	:		List[str]         =				False
   __A	:		Any         =				False






   def       __UpperCamelCase     (      self				:     List[str]  )					->						List[Any]:
       lowercase__ :   Union[str, Any]    					=				FlaxBigBirdModelTester(self  )






   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def       __UpperCamelCase     (      self				:     Optional[int]  )					->						Dict:
       super().test_from_pretrained_save_pretrained()






   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def       __UpperCamelCase     (      self				:     List[str]  )					->						Any:
       super().test_from_pretrained_with_no_automatic_init()






   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def       __UpperCamelCase     (      self				:     Tuple  )					->						str:
       super().test_no_automatic_init()






   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def       __UpperCamelCase     (      self				:     Dict  )					->						Union[str, Any]:
       super().test_hidden_states_output()






   @slow
   def       __UpperCamelCase     (      self				:     Optional[int]  )					->						Tuple:
       for model_class_name in self.all_model_classes:
           lowercase__ :   Optional[Any]    					=				model_class_name.from_pretrained("google/bigbird-roberta-base"  )
           self.assertIsNotNone(lowercase_  )






   def       __UpperCamelCase     (      self				:     int  )					->						Optional[int]:
       if self.test_attn_probs:
           super().test_attention_outputs()






   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def       __UpperCamelCase     (      self				:     str  )					->						Any:
       lowercase__					,   lowercase__ :   Dict    					=				self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
           with self.subTest(model_class.__name__  ):
               lowercase__ :   Union[str, Any]    					=				self._prepare_for_class(lowercase_	,						lowercase_  )
               lowercase__ :   Optional[Any]    					=				model_class(lowercase_  )

               @jax.jit
               def model_jitted(lowercase_				:     Tuple	,						lowercase_				:     int=None	,						**lowercase_				:     Dict  ):
                   return model(input_ids=lowercase_	,						attention_mask=lowercase_	,						**lowercase_  )

               with self.subTest("JIT Enabled"  ):
                   lowercase__ :   int    					=				model_jitted(**lowercase_  ).to_tuple()

               with self.subTest("JIT Disabled"  ):
                   with jax.disable_jit():
                       lowercase__ :   Any    					=				model_jitted(**lowercase_  ).to_tuple()

               self.assertEqual(len(lowercase_  )	,						len(lowercase_  )  )
               for jitted_output, output in zip(lowercase_	,						lowercase_  ):
                   self.assertEqual(jitted_output.shape	,						output.shape  )







   def       __UpperCamelCase     (      self				:     List[Any]	,						lowercase_				:     str	,						lowercase_				:     Union[str, Any]	,						lowercase_				:     Optional[int]	,						lowercase_				:     List[Any]=1E-5	,						lowercase_				:     Any="outputs"	,						lowercase_				:     List[str]=None  )					->						List[Any]:
       # `bigbird_block_sparse_attention` in `FlaxBigBird` returns `attention_probs = None`, while in PyTorch version,
       # an effort was done to return `attention_probs` (yet to be verified).
       if name.startswith("outputs.attentions"  ):
           return
       else:
           super().check_pt_flax_outputs(lowercase_	,						lowercase_	,						lowercase_	,						lowercase_	,						lowercase_	,						lowercase_  )



