


'''simple docstring'''



import math






class 	UpperCamelCase__:




				def __init__(		self      :			Optional[int]				,   lowerCAmelCase      :			Union[str, Any]=0				)->					str:  # a graph with Node 0,1,...,N-1

										"""simple docstring"""




										UpperCAmelCase					    =     n
										UpperCAmelCase					    =     [
										    [math.inf for j in range(0				,   lowerCAmelCase				)] for i in range(0				,   lowerCAmelCase				)
										]  # adjacency matrix for weight
										UpperCAmelCase					    =     [
										    [math.inf for j in range(0				,   lowerCAmelCase				)] for i in range(0				,   lowerCAmelCase				)
										]  # dp[i][j] stores minimum distance from i to j




				def 					a__(		self      :			Union[str, Any]				,   lowerCAmelCase      :			List[Any]				,   lowerCAmelCase      :			Optional[int]				,   lowerCAmelCase      :			str				)->					Optional[int]:

										"""simple docstring"""




										UpperCAmelCase					    =     w




				def 					a__(		self      :			Optional[int]				)->					Optional[Any]:

										"""simple docstring"""




										for k in range(0				,   self.n				):
																for i in range(0				,   self.n				):
																						for j in range(0				,   self.n				):
																												UpperCAmelCase					    =     min(self.dp[i][j]				,   self.dp[i][k] + self.dp[k][j]				)





				def 					a__(		self      :			int				,   lowerCAmelCase      :			int				,   lowerCAmelCase      :			Dict				)->					List[str]:

										"""simple docstring"""




										return self.dp[u][v]


if __name__ == "__main__":
					_lowercase      :      Optional[int]            =					Graph(5)
					graph.add_edge(0, 2, 9)
					graph.add_edge(0, 4, 10)
					graph.add_edge(1, 3, 5)
					graph.add_edge(2, 3, 7)
					graph.add_edge(3, 0, 10)
					graph.add_edge(3, 1, 2)
					graph.add_edge(3, 2, 1)
					graph.add_edge(3, 4, 6)
					graph.add_edge(4, 1, 3)
					graph.add_edge(4, 2, 4)
					graph.add_edge(4, 3, 9)
					graph.floyd_warshall()
					graph.show_min(1, 4)
					graph.show_min(0, 3)









'''simple docstring'''



import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
	import os

	import jax.numpy as jnp
	from jax import jit

	from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
	from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

	_lowercase      :      List[str]            =					"""0.12"""  # assumed parallelism: 8


if is_torch_available():
	import torch


def 				lowerCamelCase__				(	A       :      str				,       A       :      str				,       A       :      List[Any]=None					):




						'''simple docstring'''






						if rng is None:
												UpperCAmelCase					    =     random.Random()

						UpperCAmelCase					    =     1
						for dim in shape:
												total_dims *= dim

						UpperCAmelCase					    =     []
						for _ in range(A					):
												values.append(rng.randint(0				,       vocab_size - 1					)					)

						UpperCAmelCase					    =     np.array(A				,       dtype=jnp.intaa					).reshape(A					)

						return output


def 				lowerCamelCase__				(	A       :      int				,       A       :      Optional[int]=None					):




						'''simple docstring'''






						UpperCAmelCase					    =     ids_tensor(A				,       vocab_size=2				,       rng=A					)
						# make sure that at least one token is attended to for each batch
						UpperCAmelCase					    =     1
						return attn_mask






@require_flax
class 	UpperCamelCase__:
				__magic_name__    :				Optional[int]  		=	None
				__magic_name__    :				Optional[Any]  		=	()




				def 					a__(		self      :			str				)->					Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase					    =     self.model_tester.prepare_config_and_inputs_for_common()

										# cut to half length & take max batch_size 3
										UpperCAmelCase					    =     2
										UpperCAmelCase					    =     inputs['''input_ids'''].shape[-1] // 2
										UpperCAmelCase					    =     inputs['''input_ids'''][:max_batch_size, :sequence_length]

										UpperCAmelCase					    =     jnp.ones_like(lowerCAmelCase				)
										UpperCAmelCase					    =     attention_mask[:max_batch_size, :sequence_length]

										# generate max 5 tokens
										UpperCAmelCase					    =     input_ids.shape[-1] + 5
										if config.eos_token_id is not None and config.pad_token_id is None:
																# hack to allow generate for models such as GPT2 as is done in `generate()`
																UpperCAmelCase					    =     config.eos_token_id
										return config, input_ids, attention_mask, max_length




				@is_pt_flax_cross_test
				def 					a__(		self      :			Dict				)->					Optional[int]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()
										UpperCAmelCase					    =     False
										UpperCAmelCase					    =     max_length
										UpperCAmelCase					    =     0

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model_class.__name__[4:]  # Skip the "Flax" at the beginning
																UpperCAmelCase					    =     getattr(lowerCAmelCase				,   lowerCAmelCase				)
																UpperCAmelCase					    =     pt_model_class(lowerCAmelCase				).eval()
																UpperCAmelCase					    =     load_flax_weights_in_pytorch_model(lowerCAmelCase				,   flax_model.params				)

																UpperCAmelCase					    =     flax_model.generate(lowerCAmelCase				).sequences
																UpperCAmelCase					    =     pt_model.generate(torch.tensor(lowerCAmelCase				,   dtype=torch.long				)				)

																if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
																						UpperCAmelCase					    =     flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

																self.assertListEqual(pt_generation_outputs.numpy().tolist()				,   flax_generation_outputs.tolist()				)




				def 					a__(		self      :			Any				)->					Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()
										UpperCAmelCase					    =     False
										UpperCAmelCase					    =     max_length

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)




				def 					a__(		self      :			Optional[Any]				)->					int:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()
										UpperCAmelCase					    =     True
										UpperCAmelCase					    =     max_length

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)




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

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()
										UpperCAmelCase					    =     False
										UpperCAmelCase					    =     max_length
										UpperCAmelCase					    =     2

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)




				def 					a__(		self      :			List[Any]				)->					List[str]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()
										UpperCAmelCase					    =     False
										UpperCAmelCase					    =     max_length
										UpperCAmelCase					    =     2
										UpperCAmelCase					    =     2

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[0]				,   input_ids.shape[0] * config.num_return_sequences				)




				def 					a__(		self      :			Tuple				)->					List[str]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()
										UpperCAmelCase					    =     True
										UpperCAmelCase					    =     max_length
										UpperCAmelCase					    =     0.8
										UpperCAmelCase					    =     10
										UpperCAmelCase					    =     0.3
										UpperCAmelCase					    =     1
										UpperCAmelCase					    =     8
										UpperCAmelCase					    =     9

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)




				def 					a__(		self      :			Optional[Any]				)->					Optional[int]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()
										UpperCAmelCase					    =     max_length
										UpperCAmelCase					    =     1
										UpperCAmelCase					    =     8
										UpperCAmelCase					    =     9

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)




				def 					a__(		self      :			Tuple				)->					Tuple:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()
										UpperCAmelCase					    =     max_length
										UpperCAmelCase					    =     2
										UpperCAmelCase					    =     1
										UpperCAmelCase					    =     8
										UpperCAmelCase					    =     9

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)




				def 					a__(		self      :			Union[str, Any]				)->					Any:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()

										# pad attention mask on the left
										UpperCAmelCase					    =     attention_mask.at[(0, 0)].set(0				)

										UpperCAmelCase					    =     False
										UpperCAmelCase					    =     max_length

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				,   attention_mask=lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				,   attention_mask=lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)




				def 					a__(		self      :			Optional[Any]				)->					int:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()

										# pad attention mask on the left
										UpperCAmelCase					    =     attention_mask.at[(0, 0)].set(0				)

										UpperCAmelCase					    =     True
										UpperCAmelCase					    =     max_length

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				,   attention_mask=lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				,   attention_mask=lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)





				def 					a__(		self      :			Tuple				)->					Union[str, Any]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()

										# pad attention mask on the left
										UpperCAmelCase					    =     attention_mask.at[(0, 0)].set(0				)

										UpperCAmelCase					    =     2
										UpperCAmelCase					    =     max_length

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				,   attention_mask=lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				,   attention_mask=lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)







@require_flax
class 	UpperCamelCase__(       unittest.TestCase	):




				def 					a__(		self      :			Union[str, Any]				)->					Optional[int]:

										"""simple docstring"""




										UpperCAmelCase					    =     AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-bert'''				)
										UpperCAmelCase					    =     FlaxAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''				)

										UpperCAmelCase					    =     '''Hello world'''
										UpperCAmelCase					    =     tokenizer(lowerCAmelCase				,   return_tensors='''np'''				).input_ids

										# typos are quickly detected (the correct argument is `do_sample`)
										with self.assertRaisesRegex(lowerCAmelCase				,   '''do_samples'''				):
																model.generate(lowerCAmelCase				,   do_samples=lowerCAmelCase				)

										# arbitrary arguments that will not be used anywhere are also not accepted
										with self.assertRaisesRegex(lowerCAmelCase				,   '''foo'''				):
																UpperCAmelCase					    =     {'''foo''': '''bar'''}
																model.generate(lowerCAmelCase				,   **lowerCAmelCase				)





