


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

    __UpperCamelCase    :						List[str]     =							'0.12'  # assumed parallelism: 8


if is_torch_available():
    import torch





def   _UpperCAmelCase       (     UpperCAmelCase      :			Optional[Any]						,	UpperCAmelCase      :			str						,	UpperCAmelCase      :			int=None      ):
      """simple docstring"""




      if rng is None:
            __lowerCamelCase     :      List[Any]       						=						random.Random()

      __lowerCamelCase     :      int       						=						1
      for dim in shape:
            total_dims *= dim

      __lowerCamelCase     :      List[Any]       						=						[]
      for _ in range(UpperCAmelCase      ):
            values.append(rng.randint(0						,	vocab_size - 1      )      )

      __lowerCamelCase     :      Optional[Any]       						=						np.array(UpperCAmelCase						,	dtype=jnp.intaa      ).reshape(UpperCAmelCase      )

      return output





def   _UpperCAmelCase       (     UpperCAmelCase      :			int						,	UpperCAmelCase      :			Any=None      ):
      """simple docstring"""




      __lowerCamelCase     :      Tuple       						=						ids_tensor(UpperCAmelCase						,	vocab_size=2						,	rng=UpperCAmelCase      )
      # make sure that at least one token is attended to for each batch
      __lowerCamelCase     :      int       						=						1
      return attn_mask



@require_flax
class 							_UpperCamelCase :

   '''simple docstring'''







   a_				:							str    = None
   a_				:							List[str]    = ()

   def   _snake_case    (  self  : str							):






         '''simple docstring'''



         __lowerCamelCase						, __lowerCamelCase     :      int       						=						self.model_tester.prepare_config_and_inputs_for_common()

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

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

         # generate max 5 tokens
         __lowerCamelCase     :      List[str]       						=						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     :      str       						=						config.eos_token_id
         return config, input_ids, attention_mask, max_length

   @is_pt_flax_cross_test
   def   _snake_case    (  self  : str							):






         '''simple docstring'''



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

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

               __lowerCamelCase     :      str       						=						model_class.__name__[4:]  # Skip the "Flax" at the beginning
               __lowerCamelCase     :      List[str]       						=						getattr(_lowerCamelCase       ,  _lowerCamelCase							)
               __lowerCamelCase     :      Optional[Any]       						=						pt_model_class(_lowerCamelCase							).eval()
               __lowerCamelCase     :      Union[str, Any]       						=						load_flax_weights_in_pytorch_model(_lowerCamelCase       ,  flax_model.params							)

               __lowerCamelCase     :      Dict       						=						flax_model.generate(_lowerCamelCase							).sequences
               __lowerCamelCase     :      Union[str, Any]       						=						pt_model.generate(torch.tensor(_lowerCamelCase       ,  dtype=torch.long							)							)

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

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

   def   _snake_case    (  self  : Optional[Any]							):






         '''simple docstring'''



         __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase     :      Union[str, Any]       						=						self._get_input_ids_and_config()
         __lowerCamelCase     :      Any       						=						False
         __lowerCamelCase     :      List[Any]       						=						max_length

         for model_class in self.all_generative_model_classes:
               __lowerCamelCase     :      Optional[int]       						=						model_class(_lowerCamelCase							)

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

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

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

   def   _snake_case    (  self  : str							):






         '''simple docstring'''



         __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase     :      List[Any]       						=						self._get_input_ids_and_config()
         __lowerCamelCase     :      Optional[Any]       						=						True
         __lowerCamelCase     :      List[str]       						=						max_length

         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[-1]       ,  _lowerCamelCase							)

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

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

   def   _snake_case    (  self  : str							):






         '''simple docstring'''



         __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase     :      Any       						=						self._get_input_ids_and_config()
         __lowerCamelCase     :      str       						=						False
         __lowerCamelCase     :      List[Any]       						=						max_length
         __lowerCamelCase     :      Tuple       						=						2

         for model_class in self.all_generative_model_classes:
               __lowerCamelCase     :      Dict       						=						model_class(_lowerCamelCase							)

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

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

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

   def   _snake_case    (  self  : Optional[Any]							):






         '''simple docstring'''



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

         for model_class in self.all_generative_model_classes:
               __lowerCamelCase     :      List[str]       						=						model_class(_lowerCamelCase							)

               __lowerCamelCase     :      Optional[int]       						=						model.generate(_lowerCamelCase							).sequences
               self.assertEqual(generation_outputs.shape[0]       ,  input_ids.shape[0] * config.num_return_sequences							)

   def   _snake_case    (  self  : Tuple							):






         '''simple docstring'''



         __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase     :      Tuple       						=						self._get_input_ids_and_config()
         __lowerCamelCase     :      str       						=						True
         __lowerCamelCase     :      Union[str, Any]       						=						max_length
         __lowerCamelCase     :      Dict       						=						0.8
         __lowerCamelCase     :      Optional[int]       						=						1_0
         __lowerCamelCase     :      str       						=						0.3
         __lowerCamelCase     :      str       						=						1
         __lowerCamelCase     :      Optional[Any]       						=						8
         __lowerCamelCase     :      str       						=						9

         for model_class in self.all_generative_model_classes:
               __lowerCamelCase     :      Dict       						=						model_class(_lowerCamelCase							)

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

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

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

   def   _snake_case    (  self  : Any							):






         '''simple docstring'''



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

         for model_class in self.all_generative_model_classes:
               __lowerCamelCase     :      int       						=						model_class(_lowerCamelCase							)

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

               __lowerCamelCase     :      Any       						=						jit(model.generate							)
               __lowerCamelCase     :      Dict       						=						jit_generate(_lowerCamelCase							).sequences

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

   def   _snake_case    (  self  : Dict							):






         '''simple docstring'''



         __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase     :      Dict       						=						self._get_input_ids_and_config()
         __lowerCamelCase     :      Optional[Any]       						=						max_length
         __lowerCamelCase     :      int       						=						2
         __lowerCamelCase     :      Tuple       						=						1
         __lowerCamelCase     :      List[Any]       						=						8
         __lowerCamelCase     :      Optional[Any]       						=						9

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

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

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

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

   def   _snake_case    (  self  : str							):






         '''simple docstring'''



         __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase     :      Dict       						=						self._get_input_ids_and_config()

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

         __lowerCamelCase     :      str       						=						False
         __lowerCamelCase     :      Any       						=						max_length

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

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

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

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

   def   _snake_case    (  self  : Any							):






         '''simple docstring'''



         __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase     :      int       						=						self._get_input_ids_and_config()

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

         __lowerCamelCase     :      Union[str, Any]       						=						True
         __lowerCamelCase     :      Any       						=						max_length

         for model_class in self.all_generative_model_classes:
               __lowerCamelCase     :      Optional[int]       						=						model_class(_lowerCamelCase							)

               __lowerCamelCase     :      Dict       						=						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()							)


   def   _snake_case    (  self  : Tuple							):






         '''simple docstring'''



         __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase						, __lowerCamelCase     :      List[str]       						=						self._get_input_ids_and_config()

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

         __lowerCamelCase     :      Optional[Any]       						=						2
         __lowerCamelCase     :      Optional[int]       						=						max_length

         for model_class in self.all_generative_model_classes:
               __lowerCamelCase     :      int       						=						model_class(_lowerCamelCase							)

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

               __lowerCamelCase     :      int       						=						jit(model.generate							)
               __lowerCamelCase     :      Tuple       						=						jit_generate(_lowerCamelCase       ,  attention_mask=_lowerCamelCase							).sequences

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



@require_flax
class 							_UpperCamelCase (						unittest.TestCase						):

   '''simple docstring'''







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






         '''simple docstring'''



         __lowerCamelCase     :      Dict       						=						AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-bert"""							)
         __lowerCamelCase     :      Any       						=						FlaxAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""							)

         __lowerCamelCase     :      Dict       						=						"""Hello world"""
         __lowerCamelCase     :      Optional[int]       						=						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							)







def   _UpperCAmelCase       (     UpperCAmelCase      :			list      ):
      """simple docstring"""




      __lowerCamelCase     :      Optional[int]       						=						len(UpperCAmelCase      )
      for _ in range(UpperCAmelCase      ):
            for i in range(_ % 2						,	arr_size - 1						,	2      ):
                  if arr[i + 1] < arr[i]:
                        __lowerCamelCase						, __lowerCamelCase     :      List[str]       						=						arr[i + 1], arr[i]

      return arr


if __name__ == "__main__":
    __UpperCamelCase    :						str     =							list(range(10, 0, -1))
    print(F'''Original: {arr}. Sorted: {odd_even_transposition(arr)}''')



