




"""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				,						UpperCamelCase_				,						UpperCamelCase_=13				,						UpperCamelCase_=7				,						UpperCamelCase_=True				,						UpperCamelCase_=True				,						UpperCamelCase_=True				,						UpperCamelCase_=True				,						UpperCamelCase_=99				,						UpperCamelCase_=32				,						UpperCamelCase_=5				,						UpperCamelCase_=4				,						UpperCamelCase_=37				,						UpperCamelCase_="gelu"				,						UpperCamelCase_=0.1				,						UpperCamelCase_=0.1				,						UpperCamelCase_=5_12				,						UpperCamelCase_=16				,						UpperCamelCase_=2				,						UpperCamelCase_=0.0_2				,						UpperCamelCase_=4				,						)						->       List[str]:
					__lowercase			:  Tuple								=       parent
					__lowercase			:  List[Any]								=       batch_size
					__lowercase			:  str								=       seq_length
					__lowercase			:  List[Any]								=       is_training
					__lowercase			:  Optional[int]								=       use_attention_mask
					__lowercase			:  Dict								=       use_token_type_ids
					__lowercase			:  List[str]								=       use_labels
					__lowercase			:  List[str]								=       vocab_size
					__lowercase			:  Any								=       hidden_size
					__lowercase			:  Tuple								=       num_hidden_layers
					__lowercase			:  int								=       num_attention_heads
					__lowercase			:  Tuple								=       intermediate_size
					__lowercase			:  Dict								=       hidden_act
					__lowercase			:  Dict								=       hidden_dropout_prob
					__lowercase			:  Optional[Any]								=       attention_probs_dropout_prob
					__lowercase			:  str								=       max_position_embeddings
					__lowercase			:  Union[str, Any]								=       type_vocab_size
					__lowercase			:  int								=       type_sequence_label_size
					__lowercase			:  Union[str, Any]								=       initializer_range
					__lowercase			:  Any								=       num_choices
		def 	_lowerCamelCase							(						self       )						->       List[str]:
					__lowercase			:  Tuple								=       ids_tensor([self.batch_size, self.seq_length]				,						self.vocab_size       )

					__lowercase			:  Optional[int]								=       None
					if self.use_attention_mask:
								__lowercase			:  Optional[int]								=       random_attention_mask([self.batch_size, self.seq_length]       )

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

					__lowercase			:  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=UpperCamelCase_				,						initializer_range=self.initializer_range				,						)

					return config, input_ids, token_type_ids, attention_mask
		def 	_lowerCamelCase							(						self       )						->       int:
					__lowercase			:  str								=       self.prepare_config_and_inputs()
					__lowercase			,__lowercase			,__lowercase			,__lowercase			:  Union[str, Any]								=       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   UpperCAmelCase_    (			snake_case   ,					unittest.TestCase		):
		UpperCamelCase							=(
		    (
		        FlaxAlbertModel,
		        FlaxAlbertForPreTraining,
		        FlaxAlbertForMaskedLM,
		        FlaxAlbertForMultipleChoice,
		        FlaxAlbertForQuestionAnswering,
		        FlaxAlbertForSequenceClassification,
		        FlaxAlbertForTokenClassification,
		        FlaxAlbertForQuestionAnswering,
		    )
		    if is_flax_available()
		    else ()
		)
		def 	_lowerCamelCase							(						self       )						->       int:
					__lowercase			:  Union[str, Any]								=       FlaxAlbertModelTester(self       )
		@slow
		def 	_lowerCamelCase							(						self       )						->       str:
					for model_class_name in self.all_model_classes:
								__lowercase			:  Optional[int]								=       model_class_name.from_pretrained('''albert-base-v2'''       )
								__lowercase			:  List[str]								=       model(np.ones((1, 1)       )       )
								self.assertIsNotNone(UpperCamelCase_       )


@require_flax
class   UpperCAmelCase_    (			unittest.TestCase		):
		@slow
		def 	_lowerCamelCase							(						self       )						->       Optional[int]:
					__lowercase			:  str								=       FlaxAlbertModel.from_pretrained('''albert-base-v2'''       )
					__lowercase			:  Tuple								=       np.array([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]       )
					__lowercase			:  Any								=       np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]       )
					__lowercase			:  Any								=       model(UpperCamelCase_				,						attention_mask=UpperCamelCase_       )[0]
					__lowercase			:  Optional[int]								=       (1, 11, 7_68)
					self.assertEqual(output.shape				,						UpperCamelCase_       )
					__lowercase			:  Dict								=       np.array(
					    [[[-0.6_5_1_3, 1.5_0_3_5, -0.2_7_6_6], [-0.6_5_1_5, 1.5_0_4_6, -0.2_7_8_0], [-0.6_5_1_2, 1.5_0_4_9, -0.2_7_8_4]]]       )

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







"""simple docstring"""


import heapq as hq
import math
from collections.abc import Iterator
class     _UpperCamelCase    :
							'''simple docstring'''







							def __init__(	self       ,  __a      ):
										__lowerCAmelCase							       =     str(id_      )
										__lowerCAmelCase							       =     None
										__lowerCAmelCase							       =     None
										__lowerCAmelCase							       =     []
										__lowerCAmelCase							       =     {}  # {vertex:distance}

							def __lt__(	self       ,  __a      ):
										return self.key < other.key

							def __repr__(	self      ):
										return self.id

							def      snake_case							(	self       ,  __a      ):
										self.neighbors.append(__a      )



							def      snake_case							(	self       ,  __a       ,  __a      ):
										__lowerCAmelCase							       =     weight


def _lowerCamelCase      (						_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase       ,					_UpperCamelCase ):

			'''simple docstring'''




			graph[a - 1].add_neighbor(graph[b - 1] )
			graph[b - 1].add_neighbor(graph[a - 1] )
			# add the edges:
			graph[a - 1].add_edge(graph[b - 1]       ,					_UpperCamelCase )
			graph[b - 1].add_edge(graph[a - 1]       ,					_UpperCamelCase )


def _lowerCamelCase      (						_UpperCamelCase       ,					_UpperCamelCase ):

			'''simple docstring'''




			__lowerCAmelCase							       =     []
			for u in graph:
						__lowerCAmelCase							       =     math.inf
						__lowerCAmelCase							       =     None
			__lowerCAmelCase							       =     0
			__lowerCAmelCase							       =     graph[:]
			while q:
						__lowerCAmelCase							       =     min(_UpperCamelCase )
						q.remove(_UpperCamelCase )
						for v in u.neighbors:
									if (v in q) and (u.edges[v.id] < v.key):
												__lowerCAmelCase							       =     u
												__lowerCAmelCase							       =     u.edges[v.id]
			for i in range(1       ,					len(_UpperCamelCase ) ):
						a.append((int(graph[i].id ) + 1, int(graph[i].pi.id ) + 1) )
			return a


def _lowerCamelCase      (						_UpperCamelCase       ,					_UpperCamelCase ):

			'''simple docstring'''




			for u in graph:
						__lowerCAmelCase							       =     math.inf
						__lowerCAmelCase							       =     None
			__lowerCAmelCase							       =     0

			__lowerCAmelCase							       =     list(_UpperCamelCase )
			hq.heapify(_UpperCamelCase )

			while h:
						__lowerCAmelCase							       =     hq.heappop(_UpperCamelCase )
						for v in u.neighbors:
									if (v in h) and (u.edges[v.id] < v.key):
												__lowerCAmelCase							       =     u
												__lowerCAmelCase							       =     u.edges[v.id]
												hq.heapify(_UpperCamelCase )

			for i in range(1       ,					len(_UpperCamelCase ) ):
						yield (int(graph[i].id ) + 1, int(graph[i].pi.id ) + 1)


def _lowerCamelCase      (						):

			'''simple docstring'''






if __name__ == "__main__":
	import doctest

	doctest.testmod()




