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


   def 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase  ):
        '''simple docstring'''



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


   def 	snake_case__	(     self  ):
        '''simple docstring'''



        lowercase__				 =							None
        lowercase__				 =							20

        lowercase__				 =							self._get_uniform_logits(batch_size=2,				length=_a  )

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

        # compute softmax
        lowercase__				 =							jax.nn.softmax(_a,				axis=-1  )

        lowercase__				 =							FlaxTemperatureLogitsWarper(temperature=0.5  )
        lowercase__				 =							FlaxTemperatureLogitsWarper(temperature=1.3  )

        lowercase__				 =							jax.nn.softmax(temp_dist_warper_sharper(_a,				scores.copy(),				cur_len=_a  ),				axis=-1  )
        lowercase__				 =							jax.nn.softmax(temp_dist_warper_smoother(_a,				scores.copy(),				cur_len=_a  ),				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 	snake_case__	(     self  ):
        '''simple docstring'''



        lowercase__				 =							None
        lowercase__				 =							10
        lowercase__				 =							2

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

        lowercase__				 =							FlaxTopKLogitsWarper(3  )

        lowercase__				 =							top_k_warp(_a,				_a,				cur_len=_a  )

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

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

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


   def 	snake_case__	(     self  ):
        '''simple docstring'''



        lowercase__				 =							None
        lowercase__				 =							10
        lowercase__				 =							2

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

        lowercase__				 =							FlaxTopPLogitsWarper(0.8  )
        lowercase__				 =							np.exp(top_p_warp(_a,				_a,				cur_len=_a  )  )

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

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

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

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

        # 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 	snake_case__	(     self  ):
        '''simple docstring'''



        lowercase__				 =							20
        lowercase__				 =							4
        lowercase__				 =							0

        lowercase__				 =							FlaxMinLengthLogitsProcessor(min_length=10,				eos_token_id=_a  )

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

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


   def 	snake_case__	(     self  ):
        '''simple docstring'''



        lowercase__				 =							20
        lowercase__				 =							4
        lowercase__				 =							0

        lowercase__				 =							FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_a  )

        # check that all scores are -inf except the bos_token_id score
        lowercase__				 =							ids_tensor((batch_size, 1),				vocab_size=20  )
        lowercase__				 =							1
        lowercase__				 =							self._get_uniform_logits(_a,				_a  )
        lowercase__				 =							logits_processor(_a,				_a,				cur_len=_a  )
        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
        lowercase__				 =							3
        lowercase__				 =							self._get_uniform_logits(_a,				_a  )
        lowercase__				 =							logits_processor(_a,				_a,				cur_len=_a  )
        self.assertFalse(jnp.isinf(_a  ).any()  )


   def 	snake_case__	(     self  ):
        '''simple docstring'''



        lowercase__				 =							20
        lowercase__				 =							4
        lowercase__				 =							0
        lowercase__				 =							5

        lowercase__				 =							FlaxForcedEOSTokenLogitsProcessor(max_length=_a,				eos_token_id=_a  )

        # check that all scores are -inf except the eos_token_id when max_length is reached
        lowercase__				 =							ids_tensor((batch_size, 4),				vocab_size=20  )
        lowercase__				 =							4
        lowercase__				 =							self._get_uniform_logits(_a,				_a  )
        lowercase__				 =							logits_processor(_a,				_a,				cur_len=_a  )
        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
        lowercase__				 =							3
        lowercase__				 =							self._get_uniform_logits(_a,				_a  )
        lowercase__				 =							logits_processor(_a,				_a,				cur_len=_a  )
        self.assertFalse(jnp.isinf(_a  ).any()  )


   def 	snake_case__	(     self  ):
        '''simple docstring'''



        lowercase__				 =							4
        lowercase__				 =							10
        lowercase__				 =							15
        lowercase__				 =							2
        lowercase__				 =							1
        lowercase__				 =							15

        # dummy input_ids and scores
        lowercase__				 =							ids_tensor((batch_size, sequence_length),				_a  )
        lowercase__				 =							input_ids.copy()

        lowercase__				 =							self._get_uniform_logits(_a,				_a  )
        lowercase__				 =							scores.copy()

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

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

        lowercase__				 =							10

        # no processor list
        lowercase__				 =							temp_dist_warp(_a,				_a,				cur_len=_a  )
        lowercase__				 =							top_k_warp(_a,				_a,				cur_len=_a  )
        lowercase__				 =							top_p_warp(_a,				_a,				cur_len=_a  )
        lowercase__				 =							min_dist_proc(_a,				_a,				cur_len=_a  )
        lowercase__				 =							bos_dist_proc(_a,				_a,				cur_len=_a  )
        lowercase__				 =							eos_dist_proc(_a,				_a,				cur_len=_a  )

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

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

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


   def 	snake_case__	(     self  ):
        '''simple docstring'''



        lowercase__				 =							4
        lowercase__				 =							10
        lowercase__				 =							15
        lowercase__				 =							2
        lowercase__				 =							1
        lowercase__				 =							15

        # dummy input_ids and scores
        lowercase__				 =							ids_tensor((batch_size, sequence_length),				_a  )
        lowercase__				 =							input_ids.copy()

        lowercase__				 =							self._get_uniform_logits(_a,				_a  )
        lowercase__				 =							scores.copy()

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

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

        lowercase__				 =							10

        # no processor list
        def run_no_processor_list(_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase  ):
             lowercase__				 =							temp_dist_warp(_a,				_a,				cur_len=_a  )
             lowercase__				 =							top_k_warp(_a,				_a,				cur_len=_a  )
             lowercase__				 =							top_p_warp(_a,				_a,				cur_len=_a  )
             lowercase__				 =							min_dist_proc(_a,				_a,				cur_len=_a  )
             lowercase__				 =							bos_dist_proc(_a,				_a,				cur_len=_a  )
             lowercase__				 =							eos_dist_proc(_a,				_a,				cur_len=_a  )
             return scores

        # with processor list
        def run_processor_list(_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase  ):
             lowercase__				 =							FlaxLogitsProcessorList(
                 [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]  )
             lowercase__				 =							processor(_a,				_a,				cur_len=_a  )
             return scores

        lowercase__				 =							jax.jit(_a  )
        lowercase__				 =							jax.jit(_a  )

        lowercase__				 =							jitted_run_no_processor_list(_a,				_a,				_a  )
        lowercase__				 =							jitted_run_processor_list(_a,				_a,				_a  )

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

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





"""simple docstring"""


from sympy import diff, lambdify, symbols
from sympy.functions import *  # noqa: F403







def 					__a   (		A    ,				A    ,				A = "x"    ,				A = 10**-10    ,				A = 1    ,				):






     '''simple docstring'''



     lowercase__				 =							symbols(A      )
     lowercase__				 =							lambdify(A    ,				A      )
     lowercase__				 =							lambdify(A    ,				diff(A    ,				A      )      )

     lowercase__				 =							starting_point

     while True:
          if diff_function(A      ) != 0:
               lowercase__				 =							prev_guess - multiplicity * func(A      ) / diff_function(
                   A      )
          else:
               raise ZeroDivisionError("Could not find root"      ) from None

          # Precision is checked by comparing the difference of consecutive guesses
          if abs(next_guess - prev_guess      ) < precision:
               return next_guess

          lowercase__				 =							next_guess


# Let's Execute
if __name__ == "__main__":
 # Find root of trigonometric function
 # Find value of pi
 print(F'The root of sin(x) = 0 is {newton_raphson("sin(x)", 2)}')

 # Find root of polynomial
 # Find fourth Root of 5
 print(F'The root of x**4 - 5 = 0 is {newton_raphson("x**4 -5", 0.4 +5j)}')

 # Find value of e
 print(
     "The root of log(y) - 1 = 0 is ",
     F'{newton_raphson("log(y) - 1", 2, variable="y")}',
 )

 # Exponential Roots
 print(
     "The root of exp(x) - 1 = 0 is",
     F'{newton_raphson("exp(x) - 1", 1_0, precision=0.005)}',
 )

 # Find root of cos(x)
 print(F'The root of cos(x) = 0 is {newton_raphson("cos(x)", 0)}')




