

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 					lowercase							(			unittest.TestCase     ):




				def __init__(   self ,  snake_case ,  snake_case=13 ,  snake_case=7 ,  snake_case=True ,  snake_case=True ,  snake_case=True ,  snake_case=True ,  snake_case=99 ,  snake_case=32 ,  snake_case=5 ,  snake_case=4 ,  snake_case=37 ,  snake_case="gelu" ,  snake_case=0.1 ,  snake_case=0.1 ,  snake_case=512 ,  snake_case=16 ,  snake_case=2 ,  snake_case=0.02 ,  snake_case=4 ,  ):
											snake_case_     							=       parent
											snake_case_     							=       batch_size
											snake_case_     							=       seq_length
											snake_case_     							=       is_training
											snake_case_     							=       use_attention_mask
											snake_case_     							=       use_token_type_ids
											snake_case_     							=       use_labels
											snake_case_     							=       vocab_size
											snake_case_     							=       hidden_size
											snake_case_     							=       num_hidden_layers
											snake_case_     							=       num_attention_heads
											snake_case_     							=       intermediate_size
											snake_case_     							=       hidden_act
											snake_case_     							=       hidden_dropout_prob
											snake_case_     							=       attention_probs_dropout_prob
											snake_case_     							=       max_position_embeddings
											snake_case_     							=       type_vocab_size
											snake_case_     							=       type_sequence_label_size
											snake_case_     							=       initializer_range
											snake_case_     							=       num_choices




				def        a				(   self			):
											snake_case_     							=       ids_tensor([self.batch_size, self.seq_length] ,  self.vocab_size			)

											snake_case_     							=       None
											if self.use_attention_mask:
																		snake_case_     							=       random_attention_mask([self.batch_size, self.seq_length]			)

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

											snake_case_     							=       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=lowerCAmelCase_ ,  initializer_range=self.initializer_range ,  )

											return config, input_ids, token_type_ids, attention_mask




				def        a				(   self			):
											snake_case_     							=       self.prepare_config_and_inputs()
											snake_case_							,       snake_case_							,       snake_case_							,       snake_case_     							=       config_and_inputs
											snake_case_     							=       {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': attention_mask}
											return config, inputs_dict







@require_flax
class 					lowercase							(			a_      ,    unittest.TestCase     ):
				__SCREAMING_SNAKE_CASE   :					List[str]			     =     (
				    (
				        FlaxAlbertModel,
				        FlaxAlbertForPreTraining,
				        FlaxAlbertForMaskedLM,
				        FlaxAlbertForMultipleChoice,
				        FlaxAlbertForQuestionAnswering,
				        FlaxAlbertForSequenceClassification,
				        FlaxAlbertForTokenClassification,
				        FlaxAlbertForQuestionAnswering,
				    )
				    if is_flax_available()
				    else ()
				)




				def        a				(   self			):
											snake_case_     							=       FlaxAlbertModelTester(self			)




				@slow
				def        a				(   self			):
											for model_class_name in self.all_model_classes:
																		snake_case_     							=       model_class_name.from_pretrained('albert-base-v2'			)
																		snake_case_     							=       model(np.ones((1, 1)			)			)
																		self.assertIsNotNone(lowerCAmelCase_			)







@require_flax
class 					lowercase							(			unittest.TestCase     ):




				@slow
				def        a				(   self			):
											snake_case_     							=       FlaxAlbertModel.from_pretrained('albert-base-v2'			)
											snake_case_     							=       np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]			)
											snake_case_     							=       np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]			)
											snake_case_     							=       model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_			)[0]
											snake_case_     							=       (1, 11, 768)
											self.assertEqual(output.shape ,  lowerCAmelCase_			)
											snake_case_     							=       np.array(
											    [[[-0.65_13, 1.50_35, -0.27_66], [-0.65_15, 1.50_46, -0.27_80], [-0.65_12, 1.50_49, -0.27_84]]]			)

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








'''simple docstring'''





from __future__ import annotations

from functools import lru_cache
from math import ceil

UpperCAmelCase__  			=				1_0_0

UpperCAmelCase__  			=				set(range(3, NUM_PRIMES, 2))
primes.add(2)
UpperCAmelCase__  			=				42

for prime in range(3, ceil(NUM_PRIMES**0.5), 2):
					if prime not in primes:
										continue
					primes.difference_update(set(range(prime * prime, NUM_PRIMES, prime)))






@lru_cache(maxsize=100 )
def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: int ):






			"""simple docstring"""




			if number_to_partition < 0:
						return set()
			elif number_to_partition == 0:
						return {1}

			__A=	set()
			__A=	42
			__A=	42

			for prime in primes:
						if prime > number_to_partition:
									continue
						for sub in partition(number_to_partition - prime ):
									ret.add(sub * prime )

			return ret






def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: int = 5000 ):






			"""simple docstring"""




			for number_to_partition in range(1,_SCREAMING_SNAKE_CASE ):
						if len(partition(_SCREAMING_SNAKE_CASE ) ) > number_unique_partitions:
									return number_to_partition
			return None


if __name__ == "__main__":
					print(F"""{solution() = }""")



