





'''simple docstring'''


import numpy as np
import qiskit







def    _lowerCamelCase  (__lowerCamelCase  :     int = 8					,  __lowerCamelCase  :     int | None = None     )   ->    str:
			a__     = np.random.default_rng(seed=__lowerCamelCase     )

			# Roughly 25% of the qubits will contribute to the key.
			# So we take more than we need.
			a__     = 6 * key_len
			# Measurement basis for Alice's qubits.
			a__     = rng.integers(2					,  size=__lowerCamelCase     )
			# The set of states Alice will prepare.
			a__     = rng.integers(2					,  size=__lowerCamelCase     )
			# Measurement basis for Bob's qubits.
			a__     = rng.integers(2					,  size=__lowerCamelCase     )

			# Quantum Circuit to simulate BB84
			a__     = qiskit.QuantumCircuit(__lowerCamelCase					,  name="BB84"     )

			# Alice prepares her qubits according to rules above.
			for index, _ in enumerate(__lowerCamelCase     ):
						if alice_state[index] == 1:
									bbaa_circ.x(__lowerCamelCase     )
						if alice_basis[index] == 1:
									bbaa_circ.h(__lowerCamelCase     )
			bbaa_circ.barrier()

			# Bob measures the received qubits according to rules above.
			for index, _ in enumerate(__lowerCamelCase     ):
						if bob_basis[index] == 1:
									bbaa_circ.h(__lowerCamelCase     )

			bbaa_circ.barrier()
			bbaa_circ.measure_all()

			# Simulate the quantum circuit.
			a__     = qiskit.Aer.get_backend("aer_simulator"     )
			# We only need to run one shot because the key is unique.
			# Multiple shots will produce the same key.
			a__     = qiskit.execute(__lowerCamelCase					,  __lowerCamelCase					,  shots=1					,  seed_simulator=__lowerCamelCase     )
			# Returns the result of measurement.
			a__     = job.result().get_counts(__lowerCamelCase     ).most_frequent()

			# Extracting the generated key from the simulation results.
			# Only keep measurement results where Alice and Bob chose the same basis.
			a__     = "".join(
			    [
			        result_bit
			        for alice_basis_bit, bob_basis_bit, result_bit in zip(
			            __lowerCamelCase					,  __lowerCamelCase					,  __lowerCamelCase     )
			        if alice_basis_bit == bob_basis_bit
			    ]     )

			# Get final key. Pad with 0 if too short, otherwise truncate.
			a__     = gen_key[:key_len] if len(__lowerCamelCase     ) >= key_len else gen_key.ljust(__lowerCamelCase					,  "0"     )
			return key


if __name__ == "__main__":
		print(f"""The generated key is : {bbaa(8, seed=0)}""")
		from doctest import testmod

		testmod()










'''simple docstring'''


import os
from pickle import UnpicklingError
from typing import Dict, Tuple

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict, unflatten_dict

import transformers

from .utils import logging


lowerCAmelCase_   :   Union[str, Any]       =					logging.get_logger(__name__)







def    _lowerCamelCase  (__lowerCamelCase  :     List[Any]					,  __lowerCamelCase  :     Optional[int]					,  __lowerCamelCase  :     List[str]					,  __lowerCamelCase  :     Union[str, Any]=False     )   ->    Dict:
			try:
						import torch  # noqa: F401
			except ImportError:
						logger.error(
						    "Loading a PyTorch model in Flax, requires both PyTorch and Flax to be installed. Please see"
						    " https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"
						    " instructions."     )
						raise

			if not is_sharded:
						a__     = os.path.abspath(__lowerCamelCase     )
						logger.info(f'''Loading PyTorch weights from {pt_path}'''     )

						a__     = torch.load(__lowerCamelCase					,  map_location="cpu"     )
						logger.info(f'''PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()     ):,} parameters.'''     )

						a__     = convert_pytorch_state_dict_to_flax(__lowerCamelCase					,  __lowerCamelCase     )
			else:
						# model is sharded and pytorch_checkpoint_path already contains the list of .pt shard files
						a__     = convert_pytorch_sharded_state_dict_to_flax(__lowerCamelCase					,  __lowerCamelCase     )
			return flax_state_dict







def    _lowerCamelCase  (__lowerCamelCase  :     Tuple[str]					,  __lowerCamelCase  :     np.ndarray					,  __lowerCamelCase  :     Dict[str, jnp.ndarray]					,  __lowerCamelCase  :     str					,  )   ->    (Tuple[str], np.ndarray):

			def is_key_or_prefix_key_in_dict(__lowerCamelCase  :     Tuple[str]     ) -> bool:
						return len(set(__lowerCamelCase     ) & {key, (model_prefix,) + key}     ) > 0

			# layer norm
			a__     = pt_tuple_key[:-1] + ("scale",)
			if pt_tuple_key[-1] in ["weight", "gamma"] and is_key_or_prefix_key_in_dict(__lowerCamelCase     ):
						return renamed_pt_tuple_key, pt_tensor

			# batch norm layer mean
			a__     = pt_tuple_key[:-1] + ("mean",)
			if pt_tuple_key[-1] == "running_mean" and not is_key_or_prefix_key_in_dict(__lowerCamelCase     ):
						return renamed_pt_tuple_key, pt_tensor

			# batch norm layer var
			a__     = pt_tuple_key[:-1] + ("var",)
			if pt_tuple_key[-1] == "running_var" and not is_key_or_prefix_key_in_dict(__lowerCamelCase     ):
						return renamed_pt_tuple_key, pt_tensor

			# embedding
			a__     = pt_tuple_key[:-1] + ("embedding",)
			if pt_tuple_key[-1] == "weight" and is_key_or_prefix_key_in_dict(__lowerCamelCase     ):
						return renamed_pt_tuple_key, pt_tensor

			# conv layer
			a__     = pt_tuple_key[:-1] + ("kernel",)
			if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4 and not is_key_or_prefix_key_in_dict(__lowerCamelCase     ):
						a__     = pt_tensor.transpose(2					,  3					,  1					,  0     )
						return renamed_pt_tuple_key, pt_tensor

			# linear layer
			a__     = pt_tuple_key[:-1] + ("kernel",)
			if pt_tuple_key[-1] == "weight" and not is_key_or_prefix_key_in_dict(__lowerCamelCase     ):
						a__     = pt_tensor.T
						return renamed_pt_tuple_key, pt_tensor

			# old PyTorch layer norm weight
			a__     = pt_tuple_key[:-1] + ("weight",)
			if pt_tuple_key[-1] == "gamma":
						return renamed_pt_tuple_key, pt_tensor

			# old PyTorch layer norm bias
			a__     = pt_tuple_key[:-1] + ("bias",)
			if pt_tuple_key[-1] == "beta":
						return renamed_pt_tuple_key, pt_tensor

			# New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
			a__     = None
			if pt_tuple_key[-3::2] == ("parametrizations", "original0"):
						a__     = pt_tuple_key[-2] + "_g"
			elif pt_tuple_key[-3::2] == ("parametrizations", "original1"):
						a__     = pt_tuple_key[-2] + "_v"
			if name is not None:
						a__     = pt_tuple_key[:-3] + (name,)
						return renamed_pt_tuple_key, pt_tensor

			return pt_tuple_key, pt_tensor







def    _lowerCamelCase  (__lowerCamelCase  :     int					,  __lowerCamelCase  :     List[Any]     )   ->    Any:
			# convert pytorch tensor to numpy
			a__     = {k: v.numpy() for k, v in pt_state_dict.items()}

			a__     = flax_model.base_model_prefix

			# use params dict if the model contains batch norm layers
			if "params" in flax_model.params:
						a__     = flax_model.params["params"]
			else:
						a__     = flax_model.params
			a__     = flatten_dict(__lowerCamelCase     )

			# add batch_stats keys,values to dict
			if "batch_stats" in flax_model.params:
						a__     = flatten_dict(flax_model.params["batch_stats"]     )
						random_flax_state_dict.update(__lowerCamelCase     )

			a__     = {}

			a__     = (model_prefix not in flax_model_params) and (
			    model_prefix in {k.split("."     )[0] for k in pt_state_dict.keys()}
			)
			a__     = (model_prefix in flax_model_params) and (
			    model_prefix not in {k.split("."     )[0] for k in pt_state_dict.keys()}
			)

			# Need to change some parameters name to match Flax names
			for pt_key, pt_tensor in pt_state_dict.items():
						a__     = tuple(pt_key.split("."     )     )

						# remove base model prefix if necessary
						a__     = pt_tuple_key[0] == model_prefix
						if load_model_with_head_into_base_model and has_base_model_prefix:
									a__     = pt_tuple_key[1:]

						# Correctly rename weight parameters
						a__   ,				a__     = rename_key_and_reshape_tensor(
						    __lowerCamelCase					,  __lowerCamelCase					,  __lowerCamelCase					,  __lowerCamelCase     )

						# add model prefix if necessary
						a__     = (model_prefix,) + flax_key in random_flax_state_dict
						if load_base_model_into_model_with_head and require_base_model_prefix:
									a__     = (model_prefix,) + flax_key

						if flax_key in random_flax_state_dict:
									if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
												raise ValueError(
												    f'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
												    f'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'''     )

        # add batch stats if the model contains batchnorm layers
						if "batch_stats" in flax_model.params:
									if "mean" in flax_key[-1] or "var" in flax_key[-1]:
												a__     = jnp.asarray(__lowerCamelCase     )
												continue
									# remove num_batches_tracked key
									if "num_batches_tracked" in flax_key[-1]:
												flax_state_dict.pop(__lowerCamelCase					,  __lowerCamelCase     )
												continue

									# also add unexpected weight so that warning is thrown
									a__     = jnp.asarray(__lowerCamelCase     )

						else:
									# also add unexpected weight so that warning is thrown
									a__     = jnp.asarray(__lowerCamelCase     )

			return unflatten_dict(__lowerCamelCase     )







def    _lowerCamelCase  (__lowerCamelCase  :     Union[str, Any]					,  __lowerCamelCase  :     Any     )   ->    Union[str, Any]:
			import torch

			# Load the index
			a__     = {}
			for shard_file in shard_filenames:
						# load using msgpack utils
						a__     = torch.load(__lowerCamelCase     )
						a__     = {k: v.numpy() for k, v in pt_state_dict.items()}

						a__     = flax_model.base_model_prefix

						# use params dict if the model contains batch norm layers and then add batch_stats keys,values to dict
						if "batch_stats" in flax_model.params:
									a__     = flax_model.params["params"]

									a__     = flatten_dict(__lowerCamelCase     )
									random_flax_state_dict.update(flatten_dict(flax_model.params["batch_stats"]     )     )
						else:
									a__     = flax_model.params
									a__     = flatten_dict(__lowerCamelCase     )

						a__     = (model_prefix not in flax_model_params) and (
						    model_prefix in {k.split("."     )[0] for k in pt_state_dict.keys()}
						)
						a__     = (model_prefix in flax_model_params) and (
						    model_prefix not in {k.split("."     )[0] for k in pt_state_dict.keys()}
						)
						# Need to change some parameters name to match Flax names
						for pt_key, pt_tensor in pt_state_dict.items():
									a__     = tuple(pt_key.split("."     )     )

									# remove base model prefix if necessary
									a__     = pt_tuple_key[0] == model_prefix
									if load_model_with_head_into_base_model and has_base_model_prefix:
												a__     = pt_tuple_key[1:]

									# Correctly rename weight parameters
									a__   ,				a__     = rename_key_and_reshape_tensor(
									    __lowerCamelCase					,  __lowerCamelCase					,  __lowerCamelCase					,  __lowerCamelCase     )
									# add model prefix if necessary
									a__     = (model_prefix,) + flax_key in random_flax_state_dict
									if load_base_model_into_model_with_head and require_base_model_prefix:
												a__     = (model_prefix,) + flax_key

									if flax_key in random_flax_state_dict:
												if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
															raise ValueError(
															    f'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
															    f'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'''     )

            # add batch stats if the model contains batchnorm layers
									if "batch_stats" in flax_model.params:
												if "mean" in flax_key[-1]:
															a__     = jnp.asarray(__lowerCamelCase     )
															continue
												if "var" in flax_key[-1]:
															a__     = jnp.asarray(__lowerCamelCase     )
															continue
												# remove num_batches_tracked key
												if "num_batches_tracked" in flax_key[-1]:
															flax_state_dict.pop(__lowerCamelCase					,  __lowerCamelCase     )
															continue

												# also add unexpected weight so that warning is thrown
												a__     = jnp.asarray(__lowerCamelCase     )

									else:
												# also add unexpected weight so that warning is thrown
												a__     = jnp.asarray(__lowerCamelCase     )
			return unflatten_dict(__lowerCamelCase     )







def    _lowerCamelCase  (__lowerCamelCase  :     Tuple					,  __lowerCamelCase  :     Tuple     )   ->    Any:
			a__     = os.path.abspath(__lowerCamelCase     )
			logger.info(f'''Loading Flax weights from {flax_checkpoint_path}'''     )

			# import correct flax class
			a__     = getattr(__lowerCamelCase					,  "Flax" + model.__class__.__name__     )

			# load flax weight dict
			with open(__lowerCamelCase					,  "rb"     ) as state_f:
						try:
									a__     = from_bytes(__lowerCamelCase					,  state_f.read()     )
						except UnpicklingError:
									raise EnvironmentError(f'''Unable to convert {flax_checkpoint_path} to Flax deserializable object. '''     )

			return load_flax_weights_in_pytorch_model(__lowerCamelCase					,  __lowerCamelCase     )







def    _lowerCamelCase  (__lowerCamelCase  :     str					,  __lowerCamelCase  :     Tuple     )   ->    Dict:

			try:
						import torch  # noqa: F401
			except ImportError:
						logger.error(
						    "Loading a Flax weights in PyTorch, requires both PyTorch and Flax to be installed. Please see"
						    " https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"
						    " instructions."     )
						raise

			# check if we have bf16 weights
			a__     = flatten_dict(jax.tree_util.tree_map(lambda __lowerCamelCase     : x.dtype == jnp.bfloataa					,  __lowerCamelCase     )     ).values()
			if any(__lowerCamelCase     ):
						# convert all weights to fp32 if the are bf16 since torch.from_numpy can-not handle bf16
						# and bf16 is not fully supported in PT yet.
						logger.warning(
						    "Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` "
						    "before loading those in PyTorch model."     )
						a__     = jax.tree_util.tree_map(
						    lambda __lowerCamelCase     : params.astype(np.floataa     ) if params.dtype == jnp.bfloataa else params					,  __lowerCamelCase     )

			a__     = flatten_dict(__lowerCamelCase     )
			a__     = pt_model.state_dict()

			a__     = (pt_model.base_model_prefix in flax_state) and (
			    pt_model.base_model_prefix not in {k.split("."     )[0] for k in pt_model_dict.keys()}
			)
			a__     = (pt_model.base_model_prefix not in flax_state) and (
			    pt_model.base_model_prefix in {k.split("."     )[0] for k in pt_model_dict.keys()}
			)

			# keep track of unexpected & missing keys
			a__     = []
			a__     = set(pt_model_dict.keys()     )

			for flax_key_tuple, flax_tensor in flax_state_dict.items():
						a__     = flax_key_tuple[0] == pt_model.base_model_prefix
						a__     = ".".join((pt_model.base_model_prefix,) + flax_key_tuple     ) in pt_model_dict

						# adapt flax_key to prepare for loading from/to base model only
						if load_model_with_head_into_base_model and has_base_model_prefix:
									a__     = flax_key_tuple[1:]
						elif load_base_model_into_model_with_head and require_base_model_prefix:
									a__     = (pt_model.base_model_prefix,) + flax_key_tuple

						# rename flax weights to PyTorch format
						if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 4 and ".".join(__lowerCamelCase     ) not in pt_model_dict:
									# conv layer
									a__     = flax_key_tuple[:-1] + ("weight",)
									a__     = jnp.transpose(__lowerCamelCase					,  (3, 2, 0, 1)     )
						elif flax_key_tuple[-1] == "kernel" and ".".join(__lowerCamelCase     ) not in pt_model_dict:
									# linear layer
									a__     = flax_key_tuple[:-1] + ("weight",)
									a__     = flax_tensor.T
						elif flax_key_tuple[-1] in ["scale", "embedding"]:
									a__     = flax_key_tuple[:-1] + ("weight",)

						# adding batch stats from flax batch norm to pt
						elif "mean" in flax_key_tuple[-1]:
									a__     = flax_key_tuple[:-1] + ("running_mean",)
						elif "var" in flax_key_tuple[-1]:
									a__     = flax_key_tuple[:-1] + ("running_var",)

						if "batch_stats" in flax_state:
									a__     = ".".join(flax_key_tuple[1:]     )  # Remove the params/batch_stats header
						else:
									a__     = ".".join(__lowerCamelCase     )

						# We also need to look at `pt_model_dict` and see if there are keys requiring further transformation.
						a__     = {}
						# New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
						for key in pt_model_dict:
									a__     = key.split("."     )
									a__     = None
									if key_components[-3::2] == ["parametrizations", "original0"]:
												a__     = key_components[-2] + "_g"
									elif key_components[-3::2] == ["parametrizations", "original1"]:
												a__     = key_components[-2] + "_v"
									if name is not None:
												a__     = key_components[:-3] + [name]
												a__     = ".".join(__lowerCamelCase     )
												a__     = key

						if flax_key in special_pt_names:
									a__     = special_pt_names[flax_key]

						if flax_key in pt_model_dict:
									if flax_tensor.shape != pt_model_dict[flax_key].shape:
												raise ValueError(
												    f'''Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '''
												    f'''to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'''     )
									else:
												# add weight to pytorch dict
												a__     = np.asarray(__lowerCamelCase     ) if not isinstance(__lowerCamelCase					,  np.ndarray     ) else flax_tensor
												a__     = torch.from_numpy(__lowerCamelCase     )
												# remove from missing keys
												missing_keys.remove(__lowerCamelCase     )
						else:
									# weight is not expected by PyTorch model
									unexpected_keys.append(__lowerCamelCase     )

			pt_model.load_state_dict(__lowerCamelCase     )

			# re-transform missing_keys to list
			a__     = list(__lowerCamelCase     )

			if len(__lowerCamelCase     ) > 0:
						logger.warning(
						    "Some weights of the Flax model were not used when initializing the PyTorch model"
						    f''' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'''
						    f''' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'''
						    " (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"
						    f''' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'''
						    " to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"
						    " FlaxBertForSequenceClassification model)."     )
			else:
						logger.warning(f'''All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n'''     )
			if len(__lowerCamelCase     ) > 0:
						logger.warning(
						    f'''Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'''
						    f''' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'''
						    " use it for predictions and inference."     )
			else:
						logger.warning(
						    f'''All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n'''
						    "If your task is similar to the task the model of the checkpoint was trained on, "
						    f'''you can already use {pt_model.__class__.__name__} for predictions without further training.'''     )

			return pt_model



