





from collections.abc import Iterator, MutableMapping
from dataclasses import dataclass
from typing import Generic, TypeVar

lowerCAmelCase_    =    TypeVar('''KEY''')
lowerCAmelCase_    =    TypeVar('''VAL''')



@dataclass(frozen=_a,					slots=_a					)
class  __lowerCAmelCase	(      Generic[KEY, VAL]					):
			lowerCamelCase_					:					KEY
			lowerCamelCase_					:					VAL




class  __lowerCAmelCase	(      _Item					):






			def __init__(self   )					->       None:


							'''simple docstring'''



							super().__init__(__magic_name__ ,							__magic_name__   )






			def __bool__(self   )					->       bool:


							'''simple docstring'''



							return False




lowerCAmelCase_    =    _DeletedItem()



class  __lowerCAmelCase	(      MutableMapping[KEY, VAL]					):






			def __init__(self ,							__magic_name__ = 8 ,							__magic_name__ = 0.75   )					->       None:


							'''simple docstring'''



							snake_case_						:   List[Any]           =  initial_block_size
							snake_case_						:   list[_Item | None]           =  [None] * initial_block_size
							assert 0.0 < capacity_factor < 1.0
							snake_case_						:   List[str]           =  capacity_factor
							snake_case_						:   int           =  0






			def       lowerCamelCase  (self ,							__magic_name__   )					->       int:


							'''simple docstring'''



							return hash(__magic_name__   ) % len(self._buckets   )






			def       lowerCamelCase  (self ,							__magic_name__   )					->       int:


							'''simple docstring'''



							return (ind + 1) % len(self._buckets   )






			def       lowerCamelCase  (self ,							__magic_name__ ,							__magic_name__ ,							__magic_name__   )					->       bool:


							'''simple docstring'''



							snake_case_						:   Optional[int]           =  self._buckets[ind]
							if not stored:
											snake_case_						:   Optional[Any]           =  _Item(__magic_name__ ,							__magic_name__   )
											self._len += 1
											return True
							elif stored.key == key:
											snake_case_						:   List[Any]           =  _Item(__magic_name__ ,							__magic_name__   )
											return True
							else:
											return False






			def       lowerCamelCase  (self   )					->       bool:


							'''simple docstring'''



							snake_case_						:   Any           =  len(self._buckets   ) * self._capacity_factor
							return len(self   ) >= int(__magic_name__   )






			def       lowerCamelCase  (self   )					->       bool:


							'''simple docstring'''



							if len(self._buckets   ) <= self._initial_block_size:
											return False
							snake_case_						:   int           =  len(self._buckets   ) * self._capacity_factor / 2
							return len(self   ) < limit






			def       lowerCamelCase  (self ,							__magic_name__   )					->       None:


							'''simple docstring'''



							snake_case_						:   List[str]           =  self._buckets
							snake_case_						:   int           =  [None] * new_size
							snake_case_						:   Optional[int]           =  0
							for item in old_buckets:
											if item:
															self._add_item(item.key ,							item.val   )






			def       lowerCamelCase  (self   )					->       None:


							'''simple docstring'''



							self._resize(len(self._buckets   ) * 2   )






			def       lowerCamelCase  (self   )					->       None:


							'''simple docstring'''



							self._resize(len(self._buckets   ) // 2   )






			def       lowerCamelCase  (self ,							__magic_name__   )					->       Iterator[int]:


							'''simple docstring'''



							snake_case_						:   Dict           =  self._get_bucket_index(__magic_name__   )
							for _ in range(len(self._buckets   )   ):
											yield ind
											snake_case_						:   Tuple           =  self._get_next_ind(__magic_name__   )






			def       lowerCamelCase  (self ,							__magic_name__ ,							__magic_name__   )					->       None:


							'''simple docstring'''



							for ind in self._iterate_buckets(__magic_name__   ):
											if self._try_set(__magic_name__ ,							__magic_name__ ,							__magic_name__   ):
															break






			def __setitem__(self ,							__magic_name__ ,							__magic_name__   )					->       None:


							'''simple docstring'''



							if self._is_full():
											self._size_up()

							self._add_item(__magic_name__ ,							__magic_name__   )






			def __delitem__(self ,							__magic_name__   )					->       None:


							'''simple docstring'''



							for ind in self._iterate_buckets(__magic_name__   ):
											snake_case_						:   Optional[Any]           =  self._buckets[ind]
											if item is None:
															raise KeyError(__magic_name__   )
											if item is _deleted:
															continue
											if item.key == key:
															snake_case_						:   Union[str, Any]           =  _deleted
															self._len -= 1
															break
							if self._is_sparse():
											self._size_down()






			def __getitem__(self ,							__magic_name__   )					->       VAL:


							'''simple docstring'''



							for ind in self._iterate_buckets(__magic_name__   ):
											snake_case_						:   List[Any]           =  self._buckets[ind]
											if item is None:
															break
											if item is _deleted:
															continue
											if item.key == key:
															return item.val
							raise KeyError(__magic_name__   )






			def __len__(self   )					->       int:


							'''simple docstring'''



							return self._len






			def __iter__(self   )					->       Iterator[KEY]:


							'''simple docstring'''



							yield from (item.key for item in self._buckets if item)






			def __repr__(self   )					->       str:


							'''simple docstring'''



							snake_case_						:   List[str]           =  ''' ,'''.join(
							    F'''{item.key}: {item.val}''' for item in self._buckets if item   )
							return F'''HashMap({val_string})'''







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






			def __init__(self ,							__magic_name__ ,							__magic_name__=13 ,							__magic_name__=7 ,							__magic_name__=True ,							__magic_name__=True ,							__magic_name__=True ,							__magic_name__=True ,							__magic_name__=99 ,							__magic_name__=32 ,							__magic_name__=5 ,							__magic_name__=4 ,							__magic_name__=37 ,							__magic_name__="gelu" ,							__magic_name__=0.1 ,							__magic_name__=0.1 ,							__magic_name__=512 ,							__magic_name__=16 ,							__magic_name__=2 ,							__magic_name__=0.02 ,							__magic_name__=4 ,							)					->       Tuple:


							'''simple docstring'''



							snake_case_						:   Union[str, Any]           =  parent
							snake_case_						:   Optional[Any]           =  batch_size
							snake_case_						:   List[Any]           =  seq_length
							snake_case_						:   Tuple           =  is_training
							snake_case_						:   List[str]           =  use_attention_mask
							snake_case_						:   Any           =  use_token_type_ids
							snake_case_						:   Dict           =  use_labels
							snake_case_						:   Optional[Any]           =  vocab_size
							snake_case_						:   Dict           =  hidden_size
							snake_case_						:   List[Any]           =  num_hidden_layers
							snake_case_						:   Union[str, Any]           =  num_attention_heads
							snake_case_						:   Any           =  intermediate_size
							snake_case_						:   Optional[int]           =  hidden_act
							snake_case_						:   Optional[int]           =  hidden_dropout_prob
							snake_case_						:   Optional[Any]           =  attention_probs_dropout_prob
							snake_case_						:   Optional[int]           =  max_position_embeddings
							snake_case_						:   Optional[int]           =  type_vocab_size
							snake_case_						:   List[Any]           =  type_sequence_label_size
							snake_case_						:   Dict           =  initializer_range
							snake_case_						:   Dict           =  num_choices






			def       lowerCamelCase  (self   )					->       Tuple:


							'''simple docstring'''



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

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

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

							snake_case_						:   List[Any]           =  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=__magic_name__ ,							initializer_range=self.initializer_range ,							)

							return config, input_ids, token_type_ids, attention_mask






			def       lowerCamelCase  (self   )					->       List[str]:


							'''simple docstring'''



							snake_case_						:   Tuple           =  self.prepare_config_and_inputs()
							snake_case_      ,							snake_case_      ,							snake_case_      ,							snake_case_						:   Optional[int]           =  config_and_inputs
							snake_case_						:   int           =  {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
							return config, inputs_dict




@require_flax
class  __lowerCAmelCase	(      _a,					unittest.TestCase					):
			lowerCamelCase_					:					Optional[int]	      =				(
			    (
			        FlaxAlbertModel,
			        FlaxAlbertForPreTraining,
			        FlaxAlbertForMaskedLM,
			        FlaxAlbertForMultipleChoice,
			        FlaxAlbertForQuestionAnswering,
			        FlaxAlbertForSequenceClassification,
			        FlaxAlbertForTokenClassification,
			        FlaxAlbertForQuestionAnswering,
			    )
			    if is_flax_available()
			    else ()
			)






			def       lowerCamelCase  (self   )					->       List[str]:


							'''simple docstring'''



							snake_case_						:   Optional[Any]           =  FlaxAlbertModelTester(self   )






			@slow
			def       lowerCamelCase  (self   )					->       Tuple:


							'''simple docstring'''



							for model_class_name in self.all_model_classes:
											snake_case_						:   Dict           =  model_class_name.from_pretrained('''albert-base-v2'''   )
											snake_case_						:   Union[str, Any]           =  model(np.ones((1, 1)   )   )
											self.assertIsNotNone(__magic_name__   )




@require_flax
class  __lowerCAmelCase	(      unittest.TestCase					):






			@slow
			def       lowerCamelCase  (self   )					->       Optional[Any]:


							'''simple docstring'''



							snake_case_						:   Optional[Any]           =  FlaxAlbertModel.from_pretrained('''albert-base-v2'''   )
							snake_case_						:   Optional[int]           =  np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]   )
							snake_case_						:   Dict           =  np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]   )
							snake_case_						:   Union[str, Any]           =  model(__magic_name__ ,							attention_mask=__magic_name__   )[0]
							snake_case_						:   Tuple           =  (1, 11, 768)
							self.assertEqual(output.shape ,							__magic_name__   )
							snake_case_						:   str           =  np.array(
							    [[[-0.6_513, 1.5_035, -0.2_766], [-0.6_515, 1.5_046, -0.2_780], [-0.6_512, 1.5_049, -0.2_784]]]   )

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