


"""simple docstring"""



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







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

       '''simple docstring'''




       snake_case					:   int         =					jnp.ones((batch_size, length)		) / length
       return scores







   def 				lowerCamelCase	(			self		)			->	Optional[int]:

       '''simple docstring'''




       snake_case					:   str         =					None
       snake_case					:   str         =					20

       snake_case					:   Tuple         =					self._get_uniform_logits(batch_size=2		,			length=UpperCamelCase__		)

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

       # compute softmax
       snake_case					:   int         =					jax.nn.softmax(UpperCamelCase__		,			axis=-1		)

       snake_case					:   str         =					FlaxTemperatureLogitsWarper(temperature=0.5		)
       snake_case					:   Union[str, Any]         =					FlaxTemperatureLogitsWarper(temperature=1.3		)

       snake_case					:   List[str]         =					jax.nn.softmax(temp_dist_warper_sharper(UpperCamelCase__		,			scores.copy()		,			cur_len=UpperCamelCase__		)		,			axis=-1		)
       snake_case					:   Union[str, Any]         =					jax.nn.softmax(temp_dist_warper_smoother(UpperCamelCase__		,			scores.copy()		,			cur_len=UpperCamelCase__		)		,			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 				lowerCamelCase	(			self		)			->	Tuple:

       '''simple docstring'''




       snake_case					:   int         =					None
       snake_case					:   int         =					10
       snake_case					:   Optional[int]         =					2

       # create ramp distribution
       snake_case					:   Tuple         =					np.broadcast_to(np.arange(UpperCamelCase__		)[None, :]		,			(batch_size, vocab_size)		).copy()
       snake_case					:   Optional[int]         =					ramp_logits[1:, : vocab_size // 2] + vocab_size

       snake_case					:   Dict         =					FlaxTopKLogitsWarper(3		)

       snake_case					:   int         =					top_k_warp(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)

       # 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
       snake_case					:   Any         =					5
       snake_case					:   int         =					FlaxTopKLogitsWarper(top_k=1		,			filter_value=0.0		,			min_tokens_to_keep=3		)

       snake_case					:   str         =					np.broadcast_to(np.arange(UpperCamelCase__		)[None, :]		,			(batch_size, length)		).copy()
       snake_case					:   Any         =					top_k_warp_safety_check(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)

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







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

       '''simple docstring'''




       snake_case					:   Optional[Any]         =					None
       snake_case					:   Tuple         =					10
       snake_case					:   Union[str, Any]         =					2

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

       snake_case					:   str         =					FlaxTopPLogitsWarper(0.8		)
       snake_case					:   Dict         =					np.exp(top_p_warp(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)		)

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

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

       # make ramp_logits more extreme
       snake_case					:   List[str]         =					ramp_logits[1] * 100.0

       # make sure at least 2 tokens are kept
       snake_case					:   Optional[int]         =					FlaxTopPLogitsWarper(0.9		,			min_tokens_to_keep=2		,			filter_value=0.0		)
       snake_case					:   Optional[Any]         =					top_p_warp(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)

       # 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 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   List[str]         =					20
       snake_case					:   List[str]         =					4
       snake_case					:   str         =					0

       snake_case					:   List[Any]         =					FlaxMinLengthLogitsProcessor(min_length=10		,			eos_token_id=UpperCamelCase__		)

       # check that min length is applied at length 5
       snake_case					:   List[str]         =					ids_tensor((batch_size, 20)		,			vocab_size=20		)
       snake_case					:   Dict         =					5
       snake_case					:   Union[str, Any]         =					self._get_uniform_logits(UpperCamelCase__		,			UpperCamelCase__		)
       snake_case					:   int         =					min_dist_processor(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
       self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()		,			4 * [-float("inf"		)]		)

       # check that min length is not applied anymore at length 15
       snake_case					:   Tuple         =					self._get_uniform_logits(UpperCamelCase__		,			UpperCamelCase__		)
       snake_case					:   Optional[int]         =					15
       snake_case					:   List[str]         =					min_dist_processor(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
       self.assertFalse(jnp.isinf(UpperCamelCase__		).any()		)







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

       '''simple docstring'''




       snake_case					:   Optional[int]         =					20
       snake_case					:   Optional[int]         =					4
       snake_case					:   Dict         =					0

       snake_case					:   List[str]         =					FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCamelCase__		)

       # check that all scores are -inf except the bos_token_id score
       snake_case					:   List[str]         =					ids_tensor((batch_size, 1)		,			vocab_size=20		)
       snake_case					:   List[str]         =					1
       snake_case					:   Any         =					self._get_uniform_logits(UpperCamelCase__		,			UpperCamelCase__		)
       snake_case					:   Dict         =					logits_processor(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
       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
       snake_case					:   Dict         =					3
       snake_case					:   List[Any]         =					self._get_uniform_logits(UpperCamelCase__		,			UpperCamelCase__		)
       snake_case					:   List[str]         =					logits_processor(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
       self.assertFalse(jnp.isinf(UpperCamelCase__		).any()		)







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

       '''simple docstring'''




       snake_case					:   List[str]         =					20
       snake_case					:   Tuple         =					4
       snake_case					:   List[Any]         =					0
       snake_case					:   Any         =					5

       snake_case					:   Optional[Any]         =					FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCamelCase__		,			eos_token_id=UpperCamelCase__		)

       # check that all scores are -inf except the eos_token_id when max_length is reached
       snake_case					:   Dict         =					ids_tensor((batch_size, 4)		,			vocab_size=20		)
       snake_case					:   Any         =					4
       snake_case					:   Optional[int]         =					self._get_uniform_logits(UpperCamelCase__		,			UpperCamelCase__		)
       snake_case					:   List[Any]         =					logits_processor(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
       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
       snake_case					:   Tuple         =					3
       snake_case					:   List[str]         =					self._get_uniform_logits(UpperCamelCase__		,			UpperCamelCase__		)
       snake_case					:   List[Any]         =					logits_processor(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
       self.assertFalse(jnp.isinf(UpperCamelCase__		).any()		)







   def 				lowerCamelCase	(			self		)			->	str:

       '''simple docstring'''




       snake_case					:   Optional[Any]         =					4
       snake_case					:   Optional[Any]         =					10
       snake_case					:   List[str]         =					15
       snake_case					:   str         =					2
       snake_case					:   str         =					1
       snake_case					:   List[str]         =					15

       # dummy input_ids and scores
       snake_case					:   Optional[Any]         =					ids_tensor((batch_size, sequence_length)		,			UpperCamelCase__		)
       snake_case					:   Optional[int]         =					input_ids.copy()

       snake_case					:   Any         =					self._get_uniform_logits(UpperCamelCase__		,			UpperCamelCase__		)
       snake_case					:   int         =					scores.copy()

       # instantiate all dist processors
       snake_case					:   int         =					FlaxTemperatureLogitsWarper(temperature=0.5		)
       snake_case					:   str         =					FlaxTopKLogitsWarper(3		)
       snake_case					:   Optional[Any]         =					FlaxTopPLogitsWarper(0.8		)

       # instantiate all logits processors
       snake_case					:   Tuple         =					FlaxMinLengthLogitsProcessor(min_length=10		,			eos_token_id=UpperCamelCase__		)
       snake_case					:   List[str]         =					FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCamelCase__		)
       snake_case					:   Union[str, Any]         =					FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCamelCase__		,			eos_token_id=UpperCamelCase__		)

       snake_case					:   Optional[int]         =					10

       # no processor list
       snake_case					:   Optional[int]         =					temp_dist_warp(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
       snake_case					:   List[Any]         =					top_k_warp(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
       snake_case					:   Dict         =					top_p_warp(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
       snake_case					:   List[Any]         =					min_dist_proc(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
       snake_case					:   List[str]         =					bos_dist_proc(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
       snake_case					:   Dict         =					eos_dist_proc(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)

       # with processor list
       snake_case					:   Optional[int]         =					FlaxLogitsProcessorList(
           [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]		)
       snake_case					:   List[Any]         =					processor(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)

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

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







   def 				lowerCamelCase	(			self		)			->	Optional[int]:

       '''simple docstring'''




       snake_case					:   Any         =					4
       snake_case					:   Any         =					10
       snake_case					:   List[Any]         =					15
       snake_case					:   Tuple         =					2
       snake_case					:   List[str]         =					1
       snake_case					:   Dict         =					15

       # dummy input_ids and scores
       snake_case					:   Tuple         =					ids_tensor((batch_size, sequence_length)		,			UpperCamelCase__		)
       snake_case					:   List[Any]         =					input_ids.copy()

       snake_case					:   Union[str, Any]         =					self._get_uniform_logits(UpperCamelCase__		,			UpperCamelCase__		)
       snake_case					:   Dict         =					scores.copy()

       # instantiate all dist processors
       snake_case					:   List[str]         =					FlaxTemperatureLogitsWarper(temperature=0.5		)
       snake_case					:   Union[str, Any]         =					FlaxTopKLogitsWarper(3		)
       snake_case					:   Union[str, Any]         =					FlaxTopPLogitsWarper(0.8		)

       # instantiate all logits processors
       snake_case					:   str         =					FlaxMinLengthLogitsProcessor(min_length=10		,			eos_token_id=UpperCamelCase__		)
       snake_case					:   Union[str, Any]         =					FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCamelCase__		)
       snake_case					:   Any         =					FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCamelCase__		,			eos_token_id=UpperCamelCase__		)

       snake_case					:   List[str]         =					10

       # no processor list
       def run_no_processor_list(UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		):
           snake_case					:   Dict         =					temp_dist_warp(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
           snake_case					:   List[Any]         =					top_k_warp(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
           snake_case					:   Union[str, Any]         =					top_p_warp(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
           snake_case					:   Tuple         =					min_dist_proc(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
           snake_case					:   List[str]         =					bos_dist_proc(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
           snake_case					:   Optional[Any]         =					eos_dist_proc(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
           return scores

       # with processor list
       def run_processor_list(UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		):
           snake_case					:   int         =					FlaxLogitsProcessorList(
               [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]		)
           snake_case					:   Any         =					processor(UpperCamelCase__		,			UpperCamelCase__		,			cur_len=UpperCamelCase__		)
           return scores

       snake_case					:   Dict         =					jax.jit(UpperCamelCase__		)
       snake_case					:   Dict         =					jax.jit(UpperCamelCase__		)

       snake_case					:   int         =					jitted_run_no_processor_list(UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		)
       snake_case					:   List[str]         =					jitted_run_processor_list(UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		)

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

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









"""simple docstring"""



import numpy as np







class _lowerCAmelCase       :







   def __init__(			self		)			->	int:

       '''simple docstring'''




       snake_case					:   Optional[int]         =					(0, 0)
       snake_case					:   str         =					None
       snake_case					:   int         =					0
       snake_case					:   Optional[Any]         =					0
       snake_case					:   Tuple         =					0







   def __eq__(			self		,			UpperCamelCase__		)			->	Dict:

       '''simple docstring'''




       return self.position == cell.position







   def 				lowerCamelCase	(			self		)			->	str:

       '''simple docstring'''




       print(self.position		)







class _lowerCAmelCase       :







   def __init__(			self		,			UpperCamelCase__=(5, 5)		)			->	List[Any]:

       '''simple docstring'''




       snake_case					:   Optional[Any]         =					np.zeros(UpperCamelCase__		)
       snake_case					:   str         =					world_size[0]
       snake_case					:   str         =					world_size[1]







   def 				lowerCamelCase	(			self		)			->	int:

       '''simple docstring'''




       print(self.w		)







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

       '''simple docstring'''




       snake_case					:   Optional[int]         =					[
           (-1, -1),
           (-1, 0),
           (-1, 1),
           (0, -1),
           (0, 1),
           (1, -1),
           (1, 0),
           (1, 1),
       ]
       snake_case					:   Tuple         =					cell.position[0]
       snake_case					:   str         =					cell.position[1]
       snake_case					:   List[Any]         =					[]
       for n in neughbour_cord:
           snake_case					:   Any         =					current_x + n[0]
           snake_case					:   List[str]         =					current_y + n[1]
           if 0 <= x < self.world_x_limit and 0 <= y < self.world_y_limit:
               snake_case					:   Tuple         =					Cell()
               snake_case					:   int         =					(x, y)
               snake_case					:   List[Any]         =					cell
               neighbours.append(UpperCamelCase__		)
       return neighbours







def 						__lowerCAmelCase		(		lowercase						:						Union[str, Any]		,							lowercase						:						Union[str, Any]		,							lowercase						:						Union[str, Any]				)			->		Any:


    """simple docstring"""





    snake_case					:   Tuple         =					[]
    snake_case					:   List[Any]         =					[]
    _open.append(lowercase				)

    while _open:
        snake_case					:   Union[str, Any]         =					np.argmin([n.f for n in _open]				)
        snake_case					:   Tuple         =					_open[min_f]
        _closed.append(_open.pop(lowercase				)				)
        if current == goal:
            break
        for n in world.get_neigbours(lowercase				):
            for c in _closed:
                if c == n:
                    continue
            snake_case					:   int         =					current.g + 1
            snake_case      ,snake_case					:   Any         =					n.position
            snake_case      ,snake_case					:   Any         =					goal.position
            snake_case					:   int         =					(ya - ya) ** 2 + (xa - xa) ** 2
            snake_case					:   List[Any]         =					n.h + n.g

            for c in _open:
                if c == n and c.f < n.f:
                    continue
            _open.append(lowercase				)
    snake_case					:   str         =					[]
    while current.parent is not None:
        path.append(current.position				)
        snake_case					:   List[str]         =					current.parent
    path.append(current.position				)
    return path[::-1]


if __name__ == "__main__":
  __snake_case         =					Gridworld()
  # Start position and goal
  __snake_case         =					Cell()
  __snake_case         =					(0, 0)
  __snake_case         =					Cell()
  __snake_case         =					(4, 4)
  print(F'''path from {start.position} to {goal.position}''')
  __snake_case         =					astar(world, start, goal)
  #   Just for visual reasons.
  for i in s:
    __snake_case         =					1
  print(world.w)





