




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

      SCREAMING_SNAKE_CASE				=       "0.12"  # assumed parallelism: 8


if is_torch_available():
      import torch







def   _SCREAMING_SNAKE_CASE   (							lowercase_			, lowercase_			, lowercase_=None						)						->		Optional[Any]:
 if rng is None:
  A__       =				random.Random()

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

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

 A__       =				np.array(a_			, dtype=jnp.intaa						).reshape(a_						)

 return output







def   _SCREAMING_SNAKE_CASE   (							lowercase_			, lowercase_=None						)						->		Optional[int]:
 A__       =				ids_tensor(a_			, vocab_size=2			, rng=a_						)
 # make sure that at least one token is attended to for each batch
 A__       =				1
 return attn_mask




@require_flax
class 	UpperCAmelCase_  :
    lowercase__											=      None
    lowercase__											=      ()




    def 							__magic_name__							(    self						:       List[Any]      )      ->							List[Any]:




     '''simple docstring'''

     A__       =				self.model_tester.prepare_config_and_inputs_for_common()

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

     A__       =				jnp.ones_like(SCREAMING_SNAKE_CASE__      )
     A__       =				attention_mask[:max_batch_size, :sequence_length]

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




    @is_pt_flax_cross_test
    def 							__magic_name__							(    self						:       List[Any]      )      ->							List[str]:




     '''simple docstring'''

     A__       =				self._get_input_ids_and_config()
     A__       =				False
     A__       =				max_length
     A__       =				0

     for model_class in self.all_generative_model_classes:
      A__       =				model_class(SCREAMING_SNAKE_CASE__      )

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

      A__       =				flax_model.generate(SCREAMING_SNAKE_CASE__      ).sequences
      A__       =				pt_model.generate(torch.tensor(SCREAMING_SNAKE_CASE__	, dtype=torch.long      )      )

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

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




    def 							__magic_name__							(    self						:       List[str]      )      ->							int:




     '''simple docstring'''

     A__       =				self._get_input_ids_and_config()
     A__       =				False
     A__       =				max_length

     for model_class in self.all_generative_model_classes:
      A__       =				model_class(SCREAMING_SNAKE_CASE__      )

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

      A__       =				jit(model.generate      )
      A__       =				jit_generate(SCREAMING_SNAKE_CASE__      ).sequences

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




    def 							__magic_name__							(    self						:       Dict      )      ->							Dict:




     '''simple docstring'''

     A__       =				self._get_input_ids_and_config()
     A__       =				True
     A__       =				max_length

     for model_class in self.all_generative_model_classes:
      A__       =				model_class(SCREAMING_SNAKE_CASE__      )

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

      A__       =				jit(model.generate      )
      A__       =				jit_generate(SCREAMING_SNAKE_CASE__      ).sequences

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




    def 							__magic_name__							(    self						:       Tuple      )      ->							int:




     '''simple docstring'''

     A__       =				self._get_input_ids_and_config()
     A__       =				False
     A__       =				max_length
     A__       =				2

     for model_class in self.all_generative_model_classes:
      A__       =				model_class(SCREAMING_SNAKE_CASE__      )

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

      A__       =				jit(model.generate      )
      A__       =				jit_generate(SCREAMING_SNAKE_CASE__      ).sequences

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




    def 							__magic_name__							(    self						:       List[Any]      )      ->							int:




     '''simple docstring'''

     A__       =				self._get_input_ids_and_config()
     A__       =				False
     A__       =				max_length
     A__       =				2
     A__       =				2

     for model_class in self.all_generative_model_classes:
      A__       =				model_class(SCREAMING_SNAKE_CASE__      )

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




    def 							__magic_name__							(    self						:       Tuple      )      ->							Optional[int]:




     '''simple docstring'''

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

     for model_class in self.all_generative_model_classes:
      A__       =				model_class(SCREAMING_SNAKE_CASE__      )

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

      A__       =				jit(model.generate      )
      A__       =				jit_generate(SCREAMING_SNAKE_CASE__      ).sequences

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




    def 							__magic_name__							(    self						:       str      )      ->							str:




     '''simple docstring'''

     A__       =				self._get_input_ids_and_config()
     A__       =				max_length
     A__       =				1
     A__       =				8
     A__       =				9

     for model_class in self.all_generative_model_classes:
      A__       =				model_class(SCREAMING_SNAKE_CASE__      )

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

      A__       =				jit(model.generate      )
      A__       =				jit_generate(SCREAMING_SNAKE_CASE__      ).sequences

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




    def 							__magic_name__							(    self						:       int      )      ->							Dict:




     '''simple docstring'''

     A__       =				self._get_input_ids_and_config()
     A__       =				max_length
     A__       =				2
     A__       =				1
     A__       =				8
     A__       =				9

     for model_class in self.all_generative_model_classes:
      A__       =				model_class(SCREAMING_SNAKE_CASE__      )

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

      A__       =				jit(model.generate      )
      A__       =				jit_generate(SCREAMING_SNAKE_CASE__      ).sequences

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




    def 							__magic_name__							(    self						:       Tuple      )      ->							str:




     '''simple docstring'''

     A__       =				self._get_input_ids_and_config()

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

     A__       =				False
     A__       =				max_length

     for model_class in self.all_generative_model_classes:
      A__       =				model_class(SCREAMING_SNAKE_CASE__      )

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

      A__       =				jit(model.generate      )
      A__       =				jit_generate(SCREAMING_SNAKE_CASE__	, attention_mask=SCREAMING_SNAKE_CASE__      ).sequences

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




    def 							__magic_name__							(    self						:       Dict      )      ->							Tuple:




     '''simple docstring'''

     A__       =				self._get_input_ids_and_config()

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

     A__       =				True
     A__       =				max_length

     for model_class in self.all_generative_model_classes:
      A__       =				model_class(SCREAMING_SNAKE_CASE__      )

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

      A__       =				jit(model.generate      )
      A__       =				jit_generate(SCREAMING_SNAKE_CASE__	, attention_mask=SCREAMING_SNAKE_CASE__      ).sequences

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




    def 							__magic_name__							(    self						:       Optional[int]      )      ->							Dict:




     '''simple docstring'''

     A__       =				self._get_input_ids_and_config()

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

     A__       =				2
     A__       =				max_length

     for model_class in self.all_generative_model_classes:
      A__       =				model_class(SCREAMING_SNAKE_CASE__      )

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

      A__       =				jit(model.generate      )
      A__       =				jit_generate(SCREAMING_SNAKE_CASE__	, attention_mask=SCREAMING_SNAKE_CASE__      ).sequences

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






@require_flax
class 	UpperCAmelCase_  (						unittest.TestCase							):




    def 							__magic_name__							(    self						:       Dict      )      ->							Optional[Any]:




     '''simple docstring'''

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

     A__       =				'''Hello world'''
     A__       =				tokenizer(SCREAMING_SNAKE_CASE__	, return_tensors="np"      ).input_ids

     # typos are quickly detected (the correct argument is `do_sample`)
     with self.assertRaisesRegex(SCREAMING_SNAKE_CASE__	, "do_samples"      ):
      model.generate(SCREAMING_SNAKE_CASE__	, do_samples=SCREAMING_SNAKE_CASE__      )

     # arbitrary arguments that will not be used anywhere are also not accepted
     with self.assertRaisesRegex(SCREAMING_SNAKE_CASE__	, "foo"      ):
      A__       =				{'''foo''': '''bar'''}
      model.generate(SCREAMING_SNAKE_CASE__	, **SCREAMING_SNAKE_CASE__      )












"""simple docstring"""
import qiskit



def      __lowerCamelCase   (							a_     :		int ,   a_     :		int							)			->      qiskit.result.counts.Counts:
      __SCREAMING_SNAKE_CASE   :Tuple						=  qiskit.Aer.get_backend('''aer_simulator'''							)

      # Create a Quantum Circuit acting on the q register
      __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  qiskit.QuantumCircuit(a_ ,   a_							)

      # Apply X (NOT) Gate to Qubits 0 & 1
      circuit.x(0							)
      circuit.x(1							)

      # Map the quantum measurement to the classical bits
      circuit.measure([0, 1] ,   [0, 1]							)

      # Execute the circuit on the qasm simulator
      __SCREAMING_SNAKE_CASE   :Tuple						=  qiskit.execute(a_ ,   a_ ,   shots=10_00							)

      # Return the histogram data of the results of the experiment.
      return job.result().get_counts(a_							)


if __name__ == "__main__":
  lowerCamelCase_								=		single_qubit_measure(2, 2)
  print(f'Total count for various states are: {counts}')