

"""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

 _lowerCAmelCase       						=		'0.12'  # assumed parallelism: 8


if is_torch_available():
 import torch

def 		lowerCamelCase__		(       _lowerCamelCase			, _lowerCamelCase			, _lowerCamelCase=None ):

    '''simple docstring'''



    if rng is None:
        _lowerCAmelCase							: Dict	=     random.Random()

    _lowerCAmelCase							: Optional[Any]	=     1
    for dim in shape:
        total_dims *= dim

    _lowerCAmelCase							: Optional[int]	=     []
    for _ in range(A__ ):
        values.append(rng.randint(0			, vocab_size - 1 ) )

    _lowerCAmelCase							: List[Any]	=     np.array(A__			, dtype=jnp.intaa ).reshape(A__ )

    return output

def 		lowerCamelCase__		(       _lowerCamelCase			, _lowerCamelCase=None ):

    '''simple docstring'''



    _lowerCAmelCase							: Optional[int]	=     ids_tensor(A__			, vocab_size=2			, rng=A__ )
    # make sure that at least one token is attended to for each batch
    _lowerCAmelCase							: Optional[Any]	=     1
    return attn_mask







@require_flax
class 						__UpperCamelCase     :
  _UpperCAmelCase        =		None
  _UpperCAmelCase        =		()

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

      _lowerCAmelCase,       _lowerCAmelCase							: List[str]	=     self.model_tester.prepare_config_and_inputs_for_common()

      # cut to half length & take max batch_size 3
      _lowerCAmelCase							: Optional[int]	=     2
      _lowerCAmelCase							: List[str]	=     inputs['input_ids'].shape[-1] // 2
      _lowerCAmelCase							: Union[str, Any]	=     inputs['input_ids'][:max_batch_size, :sequence_length]

      _lowerCAmelCase							: Optional[Any]	=     jnp.ones_like(__lowerCamelCase   )
      _lowerCAmelCase							: int	=     attention_mask[:max_batch_size, :sequence_length]

      # generate max 5 tokens
      _lowerCAmelCase							: Optional[int]	=     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()`
          _lowerCAmelCase							: Tuple	=     config.eos_token_id
      return config, input_ids, attention_mask, max_length

  @is_pt_flax_cross_test
  def __lowerCamelCase		(		self   ):
      '''simple docstring'''

      _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase							: List[str]	=     self._get_input_ids_and_config()
      _lowerCAmelCase							: str	=     False
      _lowerCAmelCase							: Optional[int]	=     max_length
      _lowerCAmelCase							: Any	=     0

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase							: str	=     model_class(__lowerCamelCase   )

          _lowerCAmelCase							: int	=     model_class.__name__[4:]  # Skip the "Flax" at the beginning
          _lowerCAmelCase							: List[str]	=     getattr(__lowerCamelCase     ,__lowerCamelCase   )
          _lowerCAmelCase							: Tuple	=     pt_model_class(__lowerCamelCase   ).eval()
          _lowerCAmelCase							: Tuple	=     load_flax_weights_in_pytorch_model(__lowerCamelCase     ,flax_model.params   )

          _lowerCAmelCase							: List[Any]	=     flax_model.generate(__lowerCamelCase   ).sequences
          _lowerCAmelCase							: int	=     pt_model.generate(torch.tensor(__lowerCamelCase     ,dtype=torch.long   )   )

          if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
              _lowerCAmelCase							: Optional[int]	=     flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

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

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

      _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase							: Optional[int]	=     self._get_input_ids_and_config()
      _lowerCAmelCase							: Dict	=     False
      _lowerCAmelCase							: str	=     max_length

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase							: Optional[Any]	=     model_class(__lowerCamelCase   )

          _lowerCAmelCase							: Optional[Any]	=     model.generate(__lowerCamelCase   ).sequences
          self.assertEqual(generation_outputs.shape[-1]     ,__lowerCamelCase   )

          _lowerCAmelCase							: int	=     jit(model.generate   )
          _lowerCAmelCase							: List[str]	=     jit_generate(__lowerCamelCase   ).sequences

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

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

      _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase							: int	=     self._get_input_ids_and_config()
      _lowerCAmelCase							: str	=     True
      _lowerCAmelCase							: List[str]	=     max_length

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase							: List[Any]	=     model_class(__lowerCamelCase   )

          _lowerCAmelCase							: Tuple	=     model.generate(__lowerCamelCase   ).sequences
          self.assertEqual(generation_outputs.shape[-1]     ,__lowerCamelCase   )

          _lowerCAmelCase							: Optional[Any]	=     jit(model.generate   )
          _lowerCAmelCase							: Optional[Any]	=     jit_generate(__lowerCamelCase   ).sequences

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

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

      _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase							: Optional[Any]	=     self._get_input_ids_and_config()
      _lowerCAmelCase							: Optional[int]	=     False
      _lowerCAmelCase							: List[str]	=     max_length
      _lowerCAmelCase							: Dict	=     2

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase							: Optional[Any]	=     model_class(__lowerCamelCase   )

          _lowerCAmelCase							: str	=     model.generate(__lowerCamelCase   ).sequences
          self.assertEqual(generation_outputs.shape[-1]     ,__lowerCamelCase   )

          _lowerCAmelCase							: Optional[int]	=     jit(model.generate   )
          _lowerCAmelCase							: List[Any]	=     jit_generate(__lowerCamelCase   ).sequences

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

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

      _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase							: Tuple	=     self._get_input_ids_and_config()
      _lowerCAmelCase							: List[Any]	=     False
      _lowerCAmelCase							: Optional[Any]	=     max_length
      _lowerCAmelCase							: Optional[Any]	=     2
      _lowerCAmelCase							: int	=     2

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase							: Any	=     model_class(__lowerCamelCase   )

          _lowerCAmelCase							: List[Any]	=     model.generate(__lowerCamelCase   ).sequences
          self.assertEqual(generation_outputs.shape[0]     ,input_ids.shape[0] * config.num_return_sequences   )

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

      _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase							: Optional[Any]	=     self._get_input_ids_and_config()
      _lowerCAmelCase							: str	=     True
      _lowerCAmelCase							: Optional[int]	=     max_length
      _lowerCAmelCase							: List[str]	=     0.8
      _lowerCAmelCase							: Tuple	=     10
      _lowerCAmelCase							: Any	=     0.3
      _lowerCAmelCase							: int	=     1
      _lowerCAmelCase							: Union[str, Any]	=     8
      _lowerCAmelCase							: Optional[Any]	=     9

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase							: List[Any]	=     model_class(__lowerCamelCase   )

          _lowerCAmelCase							: Tuple	=     model.generate(__lowerCamelCase   ).sequences
          self.assertEqual(generation_outputs.shape[-1]     ,__lowerCamelCase   )

          _lowerCAmelCase							: str	=     jit(model.generate   )
          _lowerCAmelCase							: Tuple	=     jit_generate(__lowerCamelCase   ).sequences

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

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

      _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase							: Union[str, Any]	=     self._get_input_ids_and_config()
      _lowerCAmelCase							: List[str]	=     max_length
      _lowerCAmelCase							: Dict	=     1
      _lowerCAmelCase							: int	=     8
      _lowerCAmelCase							: Optional[Any]	=     9

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase							: Tuple	=     model_class(__lowerCamelCase   )

          _lowerCAmelCase							: Dict	=     model.generate(__lowerCamelCase   ).sequences
          self.assertEqual(generation_outputs.shape[-1]     ,__lowerCamelCase   )

          _lowerCAmelCase							: Optional[int]	=     jit(model.generate   )
          _lowerCAmelCase							: List[str]	=     jit_generate(__lowerCamelCase   ).sequences

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

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

      _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase							: List[Any]	=     self._get_input_ids_and_config()
      _lowerCAmelCase							: Optional[Any]	=     max_length
      _lowerCAmelCase							: List[Any]	=     2
      _lowerCAmelCase							: Union[str, Any]	=     1
      _lowerCAmelCase							: List[Any]	=     8
      _lowerCAmelCase							: Tuple	=     9

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase							: Optional[Any]	=     model_class(__lowerCamelCase   )

          _lowerCAmelCase							: Dict	=     model.generate(__lowerCamelCase   ).sequences
          self.assertEqual(generation_outputs.shape[-1]     ,__lowerCamelCase   )

          _lowerCAmelCase							: List[str]	=     jit(model.generate   )
          _lowerCAmelCase							: Optional[Any]	=     jit_generate(__lowerCamelCase   ).sequences

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

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

      _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase							: List[Any]	=     self._get_input_ids_and_config()

      # pad attention mask on the left
      _lowerCAmelCase							: Optional[int]	=     attention_mask.at[(0, 0)].set(0   )

      _lowerCAmelCase							: int	=     False
      _lowerCAmelCase							: List[Any]	=     max_length

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase							: Union[str, Any]	=     model_class(__lowerCamelCase   )

          _lowerCAmelCase							: Tuple	=     model.generate(__lowerCamelCase     ,attention_mask=__lowerCamelCase   ).sequences
          self.assertEqual(generation_outputs.shape[-1]     ,__lowerCamelCase   )

          _lowerCAmelCase							: List[str]	=     jit(model.generate   )
          _lowerCAmelCase							: Optional[int]	=     jit_generate(__lowerCamelCase     ,attention_mask=__lowerCamelCase   ).sequences

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

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

      _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase							: Union[str, Any]	=     self._get_input_ids_and_config()

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

      _lowerCAmelCase							: List[Any]	=     True
      _lowerCAmelCase							: Optional[int]	=     max_length

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase							: Any	=     model_class(__lowerCamelCase   )

          _lowerCAmelCase							: str	=     model.generate(__lowerCamelCase     ,attention_mask=__lowerCamelCase   ).sequences
          self.assertEqual(generation_outputs.shape[-1]     ,__lowerCamelCase   )

          _lowerCAmelCase							: int	=     jit(model.generate   )
          _lowerCAmelCase							: List[Any]	=     jit_generate(__lowerCamelCase     ,attention_mask=__lowerCamelCase   ).sequences

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







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

      _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase,       _lowerCAmelCase							: Optional[Any]	=     self._get_input_ids_and_config()

      # pad attention mask on the left
      _lowerCAmelCase							: List[Any]	=     attention_mask.at[(0, 0)].set(0   )

      _lowerCAmelCase							: Union[str, Any]	=     2
      _lowerCAmelCase							: List[Any]	=     max_length

      for model_class in self.all_generative_model_classes:
          _lowerCAmelCase							: List[Any]	=     model_class(__lowerCamelCase   )

          _lowerCAmelCase							: Optional[int]	=     model.generate(__lowerCamelCase     ,attention_mask=__lowerCamelCase   ).sequences
          self.assertEqual(generation_outputs.shape[-1]     ,__lowerCamelCase   )

          _lowerCAmelCase							: List[str]	=     jit(model.generate   )
          _lowerCAmelCase							: List[Any]	=     jit_generate(__lowerCamelCase     ,attention_mask=__lowerCamelCase   ).sequences

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







@require_flax
class 						__UpperCamelCase     (     unittest.TestCase    ):

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

      _lowerCAmelCase							: Tuple	=     AutoTokenizer.from_pretrained('hf-internal-testing/tiny-bert'   )
      _lowerCAmelCase							: List[str]	=     FlaxAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-bert-flax-only'   )

      _lowerCAmelCase							: Any	=     'Hello world'
      _lowerCAmelCase							: Optional[Any]	=     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'   ):
          _lowerCAmelCase							: Union[str, Any]	=     {'foo': 'bar'}
          model.generate(__lowerCamelCase     ,**__lowerCamelCase   )








from collections.abc import Callable

import numpy as np





def 	__a     (      A__			:  Callable			,  A__			:  float			,  A__			:  float			,  A__			:  float			,  A__			:  float       ):
    SCREAMING_SNAKE_CASE										=   int(np.ceil((x_end - xa) / step_size       )       )
    SCREAMING_SNAKE_CASE										=   np.zeros((n + 1,)       )
    SCREAMING_SNAKE_CASE										=   ya
    SCREAMING_SNAKE_CASE										=   xa

    for k in range(A__       ):
        SCREAMING_SNAKE_CASE										=   y[k] + step_size * ode_func(A__			,  y[k]       )
        SCREAMING_SNAKE_CASE										=   y[k] + (
            (step_size / 2) * (ode_func(A__			,  y[k]       ) + ode_func(x + step_size			,  A__       ))
        )
        x += step_size

    return y


if __name__ == "__main__":
    import doctest

    doctest.testmod()