





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

    a        =						"0.12"  # assumed parallelism: 8


if is_torch_available():
    import torch

def 						__magic_name__       (							__UpperCAmelCase							,	__UpperCAmelCase							,	__UpperCAmelCase=None      )						->					Any:

   '''simple docstring'''







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

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

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

   __SCREAMING_SNAKE_CASE      							=       np.array(__UpperCAmelCase							,	dtype=jnp.intaa      ).reshape(__UpperCAmelCase      )

   return output

def 						__magic_name__       (							__UpperCAmelCase							,	__UpperCAmelCase=None      )						->					Dict:

   '''simple docstring'''







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







@require_flax
class __a :
   __UpperCamelCase   : int	      =	None
   __UpperCamelCase   : Dict	      =	()





   def      UpperCAmelCase__     (					self				:			Optional[Any]    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE      							=       self.model_tester.prepare_config_and_inputs_for_common()

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

      __SCREAMING_SNAKE_CASE      							=       jnp.ones_like(lowerCamelCase    )
      __SCREAMING_SNAKE_CASE      							=       attention_mask[:max_batch_size, :sequence_length]

      # generate max 5 tokens
      __SCREAMING_SNAKE_CASE      							=       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()`
         __SCREAMING_SNAKE_CASE      							=       config.eos_token_id
      return config, input_ids, attention_mask, max_length





   @is_pt_flax_cross_test
   def      UpperCAmelCase__     (					self				:			List[str]    ):

      '''simple docstring'''

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

      for model_class in self.all_generative_model_classes:
         __SCREAMING_SNAKE_CASE      							=       model_class(lowerCamelCase    )

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

         __SCREAMING_SNAKE_CASE      							=       flax_model.generate(lowerCamelCase    ).sequences
         __SCREAMING_SNAKE_CASE      							=       pt_model.generate(torch.tensor(lowerCamelCase				,dtype=torch.long    )    )

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

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





   def      UpperCAmelCase__     (					self				:			List[Any]    ):

      '''simple docstring'''

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

      for model_class in self.all_generative_model_classes:
         __SCREAMING_SNAKE_CASE      							=       model_class(lowerCamelCase    )

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

         __SCREAMING_SNAKE_CASE      							=       jit(model.generate    )
         __SCREAMING_SNAKE_CASE      							=       jit_generate(lowerCamelCase    ).sequences

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





   def      UpperCAmelCase__     (					self				:			Optional[int]    ):

      '''simple docstring'''

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

      for model_class in self.all_generative_model_classes:
         __SCREAMING_SNAKE_CASE      							=       model_class(lowerCamelCase    )

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

         __SCREAMING_SNAKE_CASE      							=       jit(model.generate    )
         __SCREAMING_SNAKE_CASE      							=       jit_generate(lowerCamelCase    ).sequences

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





   def      UpperCAmelCase__     (					self				:			int    ):

      '''simple docstring'''

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

      for model_class in self.all_generative_model_classes:
         __SCREAMING_SNAKE_CASE      							=       model_class(lowerCamelCase    )

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

         __SCREAMING_SNAKE_CASE      							=       jit(model.generate    )
         __SCREAMING_SNAKE_CASE      							=       jit_generate(lowerCamelCase    ).sequences

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





   def      UpperCAmelCase__     (					self				:			List[Any]    ):

      '''simple docstring'''

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

      for model_class in self.all_generative_model_classes:
         __SCREAMING_SNAKE_CASE      							=       model_class(lowerCamelCase    )

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





   def      UpperCAmelCase__     (					self				:			Any    ):

      '''simple docstring'''

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

      for model_class in self.all_generative_model_classes:
         __SCREAMING_SNAKE_CASE      							=       model_class(lowerCamelCase    )

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

         __SCREAMING_SNAKE_CASE      							=       jit(model.generate    )
         __SCREAMING_SNAKE_CASE      							=       jit_generate(lowerCamelCase    ).sequences

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





   def      UpperCAmelCase__     (					self				:			Union[str, Any]    ):

      '''simple docstring'''

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

      for model_class in self.all_generative_model_classes:
         __SCREAMING_SNAKE_CASE      							=       model_class(lowerCamelCase    )

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

         __SCREAMING_SNAKE_CASE      							=       jit(model.generate    )
         __SCREAMING_SNAKE_CASE      							=       jit_generate(lowerCamelCase    ).sequences

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





   def      UpperCAmelCase__     (					self				:			Optional[Any]    ):

      '''simple docstring'''

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

      for model_class in self.all_generative_model_classes:
         __SCREAMING_SNAKE_CASE      							=       model_class(lowerCamelCase    )

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

         __SCREAMING_SNAKE_CASE      							=       jit(model.generate    )
         __SCREAMING_SNAKE_CASE      							=       jit_generate(lowerCamelCase    ).sequences

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





   def      UpperCAmelCase__     (					self				:			Optional[int]    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE      							=       self._get_input_ids_and_config()

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

      __SCREAMING_SNAKE_CASE      							=       False
      __SCREAMING_SNAKE_CASE      							=       max_length

      for model_class in self.all_generative_model_classes:
         __SCREAMING_SNAKE_CASE      							=       model_class(lowerCamelCase    )

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

         __SCREAMING_SNAKE_CASE      							=       jit(model.generate    )
         __SCREAMING_SNAKE_CASE      							=       jit_generate(lowerCamelCase				,attention_mask=lowerCamelCase    ).sequences

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





   def      UpperCAmelCase__     (					self				:			Tuple    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE      							=       self._get_input_ids_and_config()

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

      __SCREAMING_SNAKE_CASE      							=       True
      __SCREAMING_SNAKE_CASE      							=       max_length

      for model_class in self.all_generative_model_classes:
         __SCREAMING_SNAKE_CASE      							=       model_class(lowerCamelCase    )

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

         __SCREAMING_SNAKE_CASE      							=       jit(model.generate    )
         __SCREAMING_SNAKE_CASE      							=       jit_generate(lowerCamelCase				,attention_mask=lowerCamelCase    ).sequences

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





   def      UpperCAmelCase__     (					self				:			str    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE      							=       self._get_input_ids_and_config()

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

      __SCREAMING_SNAKE_CASE      							=       2
      __SCREAMING_SNAKE_CASE      							=       max_length

      for model_class in self.all_generative_model_classes:
         __SCREAMING_SNAKE_CASE      							=       model_class(lowerCamelCase    )

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

         __SCREAMING_SNAKE_CASE      							=       jit(model.generate    )
         __SCREAMING_SNAKE_CASE      							=       jit_generate(lowerCamelCase				,attention_mask=lowerCamelCase    ).sequences

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







@require_flax
class __a (   unittest.TestCase			):





   def      UpperCAmelCase__     (					self				:			List[Any]    ):

      '''simple docstring'''

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

      __SCREAMING_SNAKE_CASE      							=       """Hello world"""
      __SCREAMING_SNAKE_CASE      							=       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"""    ):
         __SCREAMING_SNAKE_CASE      							=       {"""foo""": """bar"""}
         model.generate(lowerCamelCase				,**lowerCamelCase    )












'''simple docstring'''

def 						__magic_name__       (							__UpperCAmelCase      )						->					bool:

   '''simple docstring'''







   if num < 0:
      return False

   __SCREAMING_SNAKE_CASE      							=       num
   __SCREAMING_SNAKE_CASE      							=       0
   while num > 0:
      __SCREAMING_SNAKE_CASE      							=       rev_num * 10 + (num % 10)
      num //= 10

   return num_copy == rev_num


if __name__ == "__main__":
    import doctest

    doctest.testmod()





