

import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
				import jax
				import jax.numpy as jnp

				from transformers.generation import (
				    FlaxForcedBOSTokenLogitsProcessor,
				    FlaxForcedEOSTokenLogitsProcessor,
				    FlaxLogitsProcessorList,
				    FlaxMinLengthLogitsProcessor,
				    FlaxTemperatureLogitsWarper,
				    FlaxTopKLogitsWarper,
				    FlaxTopPLogitsWarper,
				)

@require_flax
class 						__UpperCamelCase    (					unittest.TestCase						):




	def       __A					(		self	:						Dict		,	lowerCAmelCase	:						int		,	lowerCAmelCase	:						int							):


		'''simple docstring'''







		UpperCAmelCase_       				=			jnp.ones((batch_size, length)							) / length
		return scores




	def       __A					(		self	:						Dict							):


		'''simple docstring'''







		UpperCAmelCase_       				=			None
		UpperCAmelCase_       				=			20

		UpperCAmelCase_       				=			self._get_uniform_logits(batch_size=2		,	length=lowerCAmelCase							)

		# tweak scores to not be uniform anymore
		UpperCAmelCase_       				=			scores.at[1, 5].set((1 / length) + 0.1							)  # peak, 1st batch
		UpperCAmelCase_       				=			scores.at[1, 10].set((1 / length) - 0.4							)  # valley, 1st batch

		# compute softmax
		UpperCAmelCase_       				=			jax.nn.softmax(lowerCAmelCase		,	axis=-1							)

		UpperCAmelCase_       				=			FlaxTemperatureLogitsWarper(temperature=0.5							)
		UpperCAmelCase_       				=			FlaxTemperatureLogitsWarper(temperature=1.3							)

		UpperCAmelCase_       				=			jax.nn.softmax(temp_dist_warper_sharper(lowerCAmelCase		,	scores.copy()		,	cur_len=lowerCAmelCase							)		,	axis=-1							)
		UpperCAmelCase_       				=			jax.nn.softmax(temp_dist_warper_smoother(lowerCAmelCase		,	scores.copy()		,	cur_len=lowerCAmelCase							)		,	axis=-1							)

		# uniform distribution stays uniform
		self.assertTrue(jnp.allclose(probs[0, :]		,	warped_prob_sharp[0, :]		,	atol=1e-3							)							)
		self.assertTrue(jnp.allclose(probs[0, :]		,	warped_prob_smooth[0, :]		,	atol=1e-3							)							)

		# sharp peaks get higher, valleys get lower
		self.assertLess(probs[1, :].max()		,	warped_prob_sharp[1, :].max()							)
		self.assertGreater(probs[1, :].min()		,	warped_prob_sharp[1, :].min()							)

		# smooth peaks get lower, valleys get higher
		self.assertGreater(probs[1, :].max()		,	warped_prob_smooth[1, :].max()							)
		self.assertLess(probs[1, :].min()		,	warped_prob_smooth[1, :].min()							)




	def       __A					(		self	:						Optional[int]							):


		'''simple docstring'''







		UpperCAmelCase_       				=			None
		UpperCAmelCase_       				=			10
		UpperCAmelCase_       				=			2

		# create ramp distribution
		UpperCAmelCase_       				=			np.broadcast_to(np.arange(lowerCAmelCase							)[None, :]		,	(batch_size, vocab_size)							).copy()
		UpperCAmelCase_       				=			ramp_logits[1:, : vocab_size // 2] + vocab_size

		UpperCAmelCase_       				=			FlaxTopKLogitsWarper(3							)

		UpperCAmelCase_       				=			top_k_warp(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)

		# check that correct tokens are filtered
		self.assertListEqual(jnp.isinf(scores[0]							).tolist()		,	7 * [True] + 3 * [False]							)
		self.assertListEqual(jnp.isinf(scores[1]							).tolist()		,	2 * [True] + 3 * [False] + 5 * [True]							)

		# check special case
		UpperCAmelCase_       				=			5
		UpperCAmelCase_       				=			FlaxTopKLogitsWarper(top_k=1		,	filter_value=0.0		,	min_tokens_to_keep=3							)

		UpperCAmelCase_       				=			np.broadcast_to(np.arange(lowerCAmelCase							)[None, :]		,	(batch_size, length)							).copy()
		UpperCAmelCase_       				=			top_k_warp_safety_check(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)

		# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
		self.assertListEqual((scores == 0.0).sum(axis=-1							).tolist()		,	[2, 2]							)




	def       __A					(		self	:						Any							):


		'''simple docstring'''







		UpperCAmelCase_       				=			None
		UpperCAmelCase_       				=			10
		UpperCAmelCase_       				=			2

		# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
		UpperCAmelCase_       				=			np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]							)							)

		UpperCAmelCase_       				=			FlaxTopPLogitsWarper(0.8							)
		UpperCAmelCase_       				=			np.exp(top_p_warp(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)							)

		# dist should be filtered to keep min num values so that sum is >= top_p
		# exp (-inf) => 0
		UpperCAmelCase_       				=			np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]							)
		self.assertTrue(np.allclose(lowerCAmelCase		,	lowerCAmelCase		,	atol=1e-3							)							)

		# check edge cases with negative and extreme logits
		UpperCAmelCase_       				=			np.broadcast_to(np.arange(lowerCAmelCase							)[None, :]		,	(batch_size, vocab_size)							).copy() - (
		    vocab_size // 2
		)

		# make ramp_logits more extreme
		UpperCAmelCase_       				=			ramp_logits[1] * 100.0

		# make sure at least 2 tokens are kept
		UpperCAmelCase_       				=			FlaxTopPLogitsWarper(0.9		,	min_tokens_to_keep=2		,	filter_value=0.0							)
		UpperCAmelCase_       				=			top_p_warp(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)

		# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
		self.assertListEqual((filtered_dist != 0.0).sum(axis=-1							).tolist()		,	[3, 2]							)




	def       __A					(		self	:						Optional[int]							):


		'''simple docstring'''







		UpperCAmelCase_       				=			20
		UpperCAmelCase_       				=			4
		UpperCAmelCase_       				=			0

		UpperCAmelCase_       				=			FlaxMinLengthLogitsProcessor(min_length=10		,	eos_token_id=lowerCAmelCase							)

		# check that min length is applied at length 5
		UpperCAmelCase_       				=			ids_tensor((batch_size, 20)		,	vocab_size=20							)
		UpperCAmelCase_       				=			5
		UpperCAmelCase_       				=			self._get_uniform_logits(lowerCAmelCase		,	lowerCAmelCase							)
		UpperCAmelCase_       				=			min_dist_processor(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
		self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()		,	4 * [-float("inf"							)]							)

		# check that min length is not applied anymore at length 15
		UpperCAmelCase_       				=			self._get_uniform_logits(lowerCAmelCase		,	lowerCAmelCase							)
		UpperCAmelCase_       				=			15
		UpperCAmelCase_       				=			min_dist_processor(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
		self.assertFalse(jnp.isinf(lowerCAmelCase							).any()							)




	def       __A					(		self	:						int							):


		'''simple docstring'''







		UpperCAmelCase_       				=			20
		UpperCAmelCase_       				=			4
		UpperCAmelCase_       				=			0

		UpperCAmelCase_       				=			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase							)

		# check that all scores are -inf except the bos_token_id score
		UpperCAmelCase_       				=			ids_tensor((batch_size, 1)		,	vocab_size=20							)
		UpperCAmelCase_       				=			1
		UpperCAmelCase_       				=			self._get_uniform_logits(lowerCAmelCase		,	lowerCAmelCase							)
		UpperCAmelCase_       				=			logits_processor(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
		self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]							).all()							)
		self.assertListEqual(scores[:, bos_token_id].tolist()		,	4 * [0]							)  # score for bos_token_id shold be zero

		# check that bos_token_id is not forced if current length is greater than 1
		UpperCAmelCase_       				=			3
		UpperCAmelCase_       				=			self._get_uniform_logits(lowerCAmelCase		,	lowerCAmelCase							)
		UpperCAmelCase_       				=			logits_processor(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
		self.assertFalse(jnp.isinf(lowerCAmelCase							).any()							)




	def       __A					(		self	:						str							):


		'''simple docstring'''







		UpperCAmelCase_       				=			20
		UpperCAmelCase_       				=			4
		UpperCAmelCase_       				=			0
		UpperCAmelCase_       				=			5

		UpperCAmelCase_       				=			FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase		,	eos_token_id=lowerCAmelCase							)

		# check that all scores are -inf except the eos_token_id when max_length is reached
		UpperCAmelCase_       				=			ids_tensor((batch_size, 4)		,	vocab_size=20							)
		UpperCAmelCase_       				=			4
		UpperCAmelCase_       				=			self._get_uniform_logits(lowerCAmelCase		,	lowerCAmelCase							)
		UpperCAmelCase_       				=			logits_processor(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
		self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]							).all()							)
		self.assertListEqual(scores[:, eos_token_id].tolist()		,	4 * [0]							)  # score for eos_token_id should be zero

		# check that eos_token_id is not forced if max_length is not reached
		UpperCAmelCase_       				=			3
		UpperCAmelCase_       				=			self._get_uniform_logits(lowerCAmelCase		,	lowerCAmelCase							)
		UpperCAmelCase_       				=			logits_processor(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
		self.assertFalse(jnp.isinf(lowerCAmelCase							).any()							)




	def       __A					(		self	:						Optional[int]							):


		'''simple docstring'''







		UpperCAmelCase_       				=			4
		UpperCAmelCase_       				=			10
		UpperCAmelCase_       				=			15
		UpperCAmelCase_       				=			2
		UpperCAmelCase_       				=			1
		UpperCAmelCase_       				=			15

		# dummy input_ids and scores
		UpperCAmelCase_       				=			ids_tensor((batch_size, sequence_length)		,	lowerCAmelCase							)
		UpperCAmelCase_       				=			input_ids.copy()

		UpperCAmelCase_       				=			self._get_uniform_logits(lowerCAmelCase		,	lowerCAmelCase							)
		UpperCAmelCase_       				=			scores.copy()

		# instantiate all dist processors
		UpperCAmelCase_       				=			FlaxTemperatureLogitsWarper(temperature=0.5							)
		UpperCAmelCase_       				=			FlaxTopKLogitsWarper(3							)
		UpperCAmelCase_       				=			FlaxTopPLogitsWarper(0.8							)

		# instantiate all logits processors
		UpperCAmelCase_       				=			FlaxMinLengthLogitsProcessor(min_length=10		,	eos_token_id=lowerCAmelCase							)
		UpperCAmelCase_       				=			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase							)
		UpperCAmelCase_       				=			FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase		,	eos_token_id=lowerCAmelCase							)

		UpperCAmelCase_       				=			10

		# no processor list
		UpperCAmelCase_       				=			temp_dist_warp(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
		UpperCAmelCase_       				=			top_k_warp(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
		UpperCAmelCase_       				=			top_p_warp(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
		UpperCAmelCase_       				=			min_dist_proc(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
		UpperCAmelCase_       				=			bos_dist_proc(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
		UpperCAmelCase_       				=			eos_dist_proc(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)

		# with processor list
		UpperCAmelCase_       				=			FlaxLogitsProcessorList(
		    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]							)
		UpperCAmelCase_       				=			processor(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)

		# scores should be equal
		self.assertTrue(jnp.allclose(lowerCAmelCase		,	lowerCAmelCase		,	atol=1e-3							)							)

		# input_ids should never be changed
		self.assertListEqual(input_ids.tolist()		,	input_ids_comp.tolist()							)




	def       __A					(		self	:						Tuple							):


		'''simple docstring'''







		UpperCAmelCase_       				=			4
		UpperCAmelCase_       				=			10
		UpperCAmelCase_       				=			15
		UpperCAmelCase_       				=			2
		UpperCAmelCase_       				=			1
		UpperCAmelCase_       				=			15

		# dummy input_ids and scores
		UpperCAmelCase_       				=			ids_tensor((batch_size, sequence_length)		,	lowerCAmelCase							)
		UpperCAmelCase_       				=			input_ids.copy()

		UpperCAmelCase_       				=			self._get_uniform_logits(lowerCAmelCase		,	lowerCAmelCase							)
		UpperCAmelCase_       				=			scores.copy()

		# instantiate all dist processors
		UpperCAmelCase_       				=			FlaxTemperatureLogitsWarper(temperature=0.5							)
		UpperCAmelCase_       				=			FlaxTopKLogitsWarper(3							)
		UpperCAmelCase_       				=			FlaxTopPLogitsWarper(0.8							)

		# instantiate all logits processors
		UpperCAmelCase_       				=			FlaxMinLengthLogitsProcessor(min_length=10		,	eos_token_id=lowerCAmelCase							)
		UpperCAmelCase_       				=			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase							)
		UpperCAmelCase_       				=			FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase		,	eos_token_id=lowerCAmelCase							)

		UpperCAmelCase_       				=			10

		# no processor list
		def run_no_processor_list(lowerCAmelCase	:						Tuple		,	lowerCAmelCase	:						Union[str, Any]		,	lowerCAmelCase	:						str							):
			UpperCAmelCase_       				=			temp_dist_warp(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
			UpperCAmelCase_       				=			top_k_warp(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
			UpperCAmelCase_       				=			top_p_warp(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
			UpperCAmelCase_       				=			min_dist_proc(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
			UpperCAmelCase_       				=			bos_dist_proc(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
			UpperCAmelCase_       				=			eos_dist_proc(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
			return scores

		# with processor list
		def run_processor_list(lowerCAmelCase	:						List[Any]		,	lowerCAmelCase	:						Union[str, Any]		,	lowerCAmelCase	:						Optional[Any]							):
			UpperCAmelCase_       				=			FlaxLogitsProcessorList(
			    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]							)
			UpperCAmelCase_       				=			processor(lowerCAmelCase		,	lowerCAmelCase		,	cur_len=lowerCAmelCase							)
			return scores

		UpperCAmelCase_       				=			jax.jit(lowerCAmelCase							)
		UpperCAmelCase_       				=			jax.jit(lowerCAmelCase							)

		UpperCAmelCase_       				=			jitted_run_no_processor_list(lowerCAmelCase		,	lowerCAmelCase		,	lowerCAmelCase							)
		UpperCAmelCase_       				=			jitted_run_processor_list(lowerCAmelCase		,	lowerCAmelCase		,	lowerCAmelCase							)

		# scores should be equal
		self.assertTrue(jnp.allclose(lowerCAmelCase		,	lowerCAmelCase		,	atol=1e-3							)							)

		# input_ids should never be changed
		self.assertListEqual(input_ids.tolist()		,	input_ids_comp.tolist()							)


from __future__ import annotations



def       __lowerCAmelCase      (						A       ,     A    ):
	print(F"Vertex\tShortest Distance from vertex {src}"    )
	for i, d in enumerate(A    ):
		print(F"{i}\t\t{d}"    )



def       __lowerCAmelCase      (						A       ,     A       ,     A    ):
	for j in range(A    ):
		UpperCAmelCase_				,   UpperCAmelCase_				,   UpperCAmelCase_       				=			(graph[j][k] for k in ["src", "dst", "weight"])
		if distance[u] != float("inf"    ) and distance[u] + w < distance[v]:
			return True
	return False



def       __lowerCAmelCase      (						A       ,     A       ,     A       ,     A    ):
	UpperCAmelCase_       				=			[float("inf"    )] * vertex_count
	UpperCAmelCase_       				=			0.0

	for _ in range(vertex_count - 1    ):
		for j in range(A    ):
			UpperCAmelCase_				,   UpperCAmelCase_				,   UpperCAmelCase_       				=			(graph[j][k] for k in ["src", "dst", "weight"])

			if distance[u] != float("inf"    ) and distance[u] + w < distance[v]:
				UpperCAmelCase_       				=			distance[u] + w

	UpperCAmelCase_       				=			check_negative_cycle(A       ,     A       ,     A    )
	if negative_cycle_exists:
		raise Exception("Negative cycle found"    )

	return distance


if __name__ == "__main__":
				import doctest

				doctest.testmod()

				_a:    List[str] 							=			int(input("""Enter number of vertices: """).strip())
				_a:    str 							=			int(input("""Enter number of edges: """).strip())

				_a:    list[dict[str, int]] 							=			[{} for _ in range(E)]

				for i in range(E):
								print("""Edge """, i + 1)
								_a							,			_a							,			_a:    Any 							=			(
								    int(x)
								    for x in input("""Enter source, destination, weight: """).strip().split(""" """)
								)
								_a:    str 							=			{"""src""": src, """dst""": dest, """weight""": weight}

				_a:    List[Any] 							=			int(input("""\nEnter shortest path source:""").strip())
				_a:    Optional[int] 							=			bellman_ford(graph, V, E, source)
				print_distance(shortest_distance, 0)